Filtrowanie metadanych w Databricks RAG przy użyciu funkcji ai_extract i ai_prep_search

Przepływ w Databricks zbudowany na funkcji SQL AI ai_prep_search już wiele robi. Wchodzi tu skonwertowany tekst dokumentów, a na wyjściu pojawiają się fragmenty wzbogacone semantycznie i kontekstowo, gotowe do użycia w Vector Search. Proces jest uporządkowany: konwertuje się dokumenty, przygotowuje fragmenty, spłaszcza je do tabeli Delta i tworzy dla niej indeks.

Jednak coś brakuje, gdy ten proces przetwarzania jest stosowany do korpusu dokumentów, w tym przypadku do zestawu instrukcji Orion Robotics. Same fragmenty są dobre – pole chunk_to_embed zawiera odpowiedni kontekst dokumentu, a chunk_to_retrieve jest na tyle czysty, że można go przekazać modelowi językowemu. Problemy pojawiają się podczas wyszukiwania, ponieważ nie ma pewnego sposobu na filtrowanie. Prośba taka jak ">znajdź mi przewodnik po oprogramowaniu sterującym Orion Motion Controller" musi być spełniona wyłącznie na podstawie podobieństwa semantycznego, ponieważ indeks nie posiada ustrukturyzowanych kolumn takich jak document_type, product_name czy target_audience.

Dla wyszukiwania opartego na agentach istnieje prawdziwa luka. Czysta wyszukiwanie semantyczne sprawdza się w demonstracjach. W środowisku produkcyjnym agent musi zmniejszyć przestrzeń wyszukiwania przed wykonaniem porównania wektorowego, a filtry metadanych dokładnie to robią. Rozwiązaniem jest dodanie ai_extract v2.1 do procesu.

Dlaczego filtry metadanych są ważne

Hybrydowa wyszukiwanie w RAG zazwyczaj oznacza połączenie wyszukiwania wektorowego z dopasowywaniem słów kluczowych, takim jak BM25. Jest to przydatne, ale dla zorganizowanych zbiorów dokumentów ważniejszy jest inny rodzaj hybrydy: wyszukiwanie wektorowe w połączeniu z filtrowaniem metadanych.

Załóżmy agenta, który już wie, że użytkownik potrzebuje podręcznika do Orion Arm. Sortowanie 10 000 fragmentów z każdego rodzaju dokumentu byłoby marnotrawstwem. Lepszym rozwiązaniem jest najpierw ograniczenie danych do istotnych dokumentów, a następnie użycie podobieństwa wektorowego do uporządkowania najlepszych wyników w tym mniejszym zbiorze. Databricks Vector Search umożliwia to dzięki kolumnom filtrującym określonym podczas tworzenia indeksu.

Prymatem jest to, że metadane muszą być czyste, spójne i mieć prawidłowy typ. Jeśli w jednym wierszu document_type zawiera wartość „User Manual”, w innym „manual”, a w trzecim nic, filtry nie działają poprawnie. W rezultacie jakość wyszukiwania jest kiepska, a na naprawę tracimy godziny, mimo że prawdziwym winowajcą są wcześniej utworzone niespójne metadane.

ai_extract v2.1 rozwiązuje ten problem. Przyjmuje ustrukturyzowany obiekt VARIANT wygenerowany przez ai_parse_document i zwraca metadane typowane, sprawdzone pod kątem określonego schematu. Pola typu enum określają stałe wartości, pola tekstowe kontrolują wydobywanie tekstu swobodnego, a tablice obejmują pola wielowartościowe. Ponieważ proces wydobywania odbywa się zgodnie ze schematem, wynik pozostaje spójny we wszystkich dokumentach, nawet jeśli ich układ się różni.

Trzy funkcje, trzy zadania

Rolę tych trzech funkcji nie pokrywa się, więc warto je rozdzielić przed przeglądaniem notatnika.

1. ai_parse_document v2.0 przekształca surowe pliki binarne (PDF, DOCX, obrazy) w ustrukturyzowany obiekt VARIANT, identyfikując bloki tekstowe, nagłówki, tabele, rysunki, ramki ograniczające oraz hierarchię sekcji. Funkcja ta wie, co strukturalnie znajduje się w dokumencie, ale nie ma pojęcia, o czym jest ten dokument.

ai_extract v2.1 bierze przetworzony VARIANT i, zgodnie ze schematem dostarczonym przez użytkownika, wydobywa metadane na poziomie dokumentu. Rozumie temat tylko dlatego, że schemat wskazuje, czego szukać. Jego wynikiem jest nowy VARIANT z polami typowanymi, takimi jak document_title, document_type, product_name i target_audience, oraz wszelkimi innymi zdefiniowanymi polami. Wersja 2.1 wprowadziła wskaźniki pewności oraz cytaty i jest wersją, którą należy używać.

ai_prep_search dzieli przetworzony VARIANT na fragmenty semantyczne gotowe do wyszukiwania, wykorzystując wskazówki strukturalne uzyskane podczas przetwarzania (przestrzenie ograniczające, typy elementów, hierarchię sekcji). Metadane na poziomie dokumentu pozostawia nietknięte, ponieważ należą one do funkcji ai_extract.

Każda funkcja ma jedną rolę, a cały proces składa się z połączenia tych funkcji.

Krok po kroku w notatniku

Notatnik pracuje z korpusem Orion Robotics, który znajduje się w volume Unity Catalog i zawiera pliki PDF, DOCX oraz obrazy. Najpierw następuje analiza tekstu:

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

Każdy plik w tym tomie jest odczytywany jako dane binarne i przekształcany w ustrukturyzowany obiekt VARIANT. Mądrze jest wyraźnie zastosować wersję 2.0, ponieważ ten schemat zawiera współrzędne prostokąta otaczającego oraz opisy figur, od których ai_prep_search zależy przy podejmowaniu decyzji o dzieleniu na fragmenty. Pozostanie przy ustawieniach domyślnych mogłoby spowodować, że Databricks zmieni zachowanie wewnątrz procesu przetwarzania.

Przykładowe uruchomienie dało następujący wynik:

Następne zapytanie wydobywa metadane z przetworzonych danych:

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

Część wyborów wymaga wyjaśnienia. Zarówno document_type, jak i target_audience to typy wyliczalne, ponieważ filtry potrzebują zamkniętego zbioru wartości. W formie tekstu wolnego mogą pojawić się jako „manual”, „User Manual” lub „technical manual”, co utrudnia dopasowanie. Opcja instructions dostarcza kontekstu – w tym przypadku chodzi o to, że pliki to dokumenty techniczne firmy Orion Robotics – dzięki czemu proces wydobywania danych jest skuteczniejszy. Ponadto funkcja ai_extract przyjmuje wartość VARIANT z ai_parse_document bez żadnej konwersji na tekst prosty.

Wynik próbnego uruchomienia:

Następnie następuje dzielenie na fragmenty:

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

ai_prep_search wykorzystuje identyczną, przetworzoną wersję dokumentu. Dokumenty są dzielone na fragmenty przechowywane w tablicy document.contents, a każdy fragment zawiera informacje chunk_id, chunk_position, chunk_to_retrieve i chunk_to_embed. Ostatnie dwa nie przedstawiają celowo tego samego tekstu.

chunk_to_embed zawiera dodatkowy kontekst, takie jak tytuł dokumentu, nagłówek sekcji oraz szczegóły strony, aby model umożliwiający wstawianie treści mógł właściwie umieścić dany fragment. chunk_to_retrieve to bardziej uporządkowany tekst, zazwyczaj ta wersja, która jest pokazywana użytkownikom lub przekazywana modelowi LLM.

Takie podział jest rozsądnym rozwiązaniem. Wiele pipeline’ów typu RAG wstawia i odpowiada na podstawie tego samego tekstu, jednak tekst bogaty w kontekst nadaje się do wstawiania treści, natomiast tekst uporządkowany lepiej nadaje się do czytania, więc lepsze jest zachowanie obu wersji.

Przykład wyniku:

Łączenie wszystkiego razem

Złożona zapytanie powstaje w momencie połączenia tych elementów:

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

Krok LATERAL variant_explode rozszerza zagnieżdżoną tablicę fragmentów na wiersze. Operacja JOIN na ścieżce dołącza metadane na poziomie dokumentu do każdego fragmentu tego dokumentu. Wynikiem jest jeden wiersz na fragment, zawierający zarówno pola do wyszukiwania, jak i kolumny metadanych podlegających filtrowaniu. Taka jest struktura przechowywana w tabeli Delta i używana do zaopatrywania indeksu Vector Search.

Wynik próbnego uruchomienia:

Tabela jest materializowana w ten sposób:

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

Gdy tabela jest już przygotowana, tworzony jest indeks Mosaic AI Vector Search Delta Sync z następującymi ustawieniami:

1. Kolumna źródłowa do embedowania: chunk_to_embed
2. Klucz główny: chunk_id
3. Pola poddawane filtrowaniu: doc_version wraz z document_type, product_name i target_audience

chunk_id pozostaje niezmienny pomiędzy poszczególnymi uruchomieniami i składa się z ID dokumentu oraz pozycji chunka. Dzięki temu incrementalne aktualizacje typu MERGE w tabeli Delta pozostawiają nienaruszone te chunki, zamiast je ponownie indeksować, co stanowi prawdziwą oszczędność przy regularnym odświeżaniu korpusu.

Korzyści podczas wyszukiwania

Korzyść ta pojawia się w momencie zadawania pytania przez użytkownika, a nie podczas ładowania. Załóżmy, że zapytanie brzmi „Znajdź instrukcję obsługi firmware dla Orion Motion Controller”. Wtedy proces wyszukiwania może przebiegać w dwóch krokach:

1. Określ odpowiednie filtry metadanych, takie jak document_type = „guide” i product_name = „Orion Motion Controller”.
2. Wbuduj zapytanie i przeprowadź wyszukiwanie wyłącznie wśród dokumentów, które spełniają te filtry.

Z brakiem funkcji ai_extract agent polega wyłącznie na semantycznej podobieństwie. W przypadku kilku podobnych dokumentów może to być w porządku. Jednak przy tysiącach plików obejmujących wiele produktów, typów i grup odbiorców staje się to niewiarygodne. Metadane oparte na enumeracji rozwiązują ten problem: agent nie musi zgadywać, czy chodzi o podręcznik, czy przewodnik, a po prostu filtrowa według kryterium document_type = „manual”, aby uzyskać dokładne wyniki.

Ograniczenia i koszty

ai_prep_search nie umożliwia kontroli nad dzieleniem tekstu na fragmenty. Nie można ustalić rozmiaru fragmentów, stopnia ich nakładania się ani algorytmu. W wielu przypadkach wykorzystania RAG do przetwarzania dokumentów jest to akceptowalne, ponieważ semantyczne dzielenie tekstu z dodanym kontekstem często przewyższa efektywność prostego dzielenia o stałym rozmiarze. Kontrola nad fragmentami rodzicowsko-dziecięcymi lub szczegółowe zarządzanie nakładaniem się wymaga jednak stworzenia specjalnej logiki.

ai_extract zwiększa koszty i opóźnienia, ponieważ każdy dokument wymaga dodatkowej wywołania modelu LLM. Delta Sync może pominąć niezmienione dokumenty, ale pierwsze pełne przetworzenie może nadal być kosztowne, dlatego rozsądnie jest przeprowadzić testy na próbce przed przetworzeniem całego korpusu.

Należy używać wersji 2.0 funkcji ai_parse_document, gdy jej wynik VARIANT trafia do ai_prep_search lub ai_extract. Ta wersja dostarcza ramki ograniczające, typy elementów oraz opisy rysunków. Surowy tekst można przekazać do ai_extract, choć wtedy tracą się korzyści wynikające z analizy strukturalnej tekstu.

Zakończenie projektu

Dodanie ai_extract v2.1 nie wymaga zbyt wiele dodatkowego kodu, ale znacząco zmienia architekturę, ponieważ fragmenty danych zawierają teraz czyste, skategoryzowane i filtrowalne metadane, a hybrydowe wyszukiwanie staje się bardziej niezawodne. Proces składa się nadal z trzech wyraźnych etapów:

1. ai_parse_document przechwytuje strukturę dokumentu.
2. ai_extract przechwytuje przydatne metadane.
3. ai_prep_search tworzy semantyczne fragmenty danych.

Taka separacja przynosi korzyści przy pracy z dokumentami różnych produktów, wersji, grup odbiorców i typów. Bez filtrów metadanych wyszukiwanie opiera się głównie na wyszukiwaniu podobieństwa. Z ich użyciem agent najpierw zawęża zakres wyników, a dopiero potem je sortuje.

Ten podejście nie sprawdzi się we wszystkich przypadkach. Metadane na poziomie fragmentów, wyszukiwanie elementów rodzic-dziecko czy dostosowane dzielenie na fragmenty wciąż wymagają specjalnego kodu. Dla wielu pipeline’ów RAG w Databricks jest to jednak czystsze i prostsze miejsce na rozpoczęcie pracy. Ponadto fragmenty powinny być nadal testowane i oceniane, a nigdy nie można im ufać ślepo.