Narzędzia AI agenta do nieustrukturyzowanego wyszukiwania

Ważny

Ta funkcja jest dostępna w publicznej wersji zapoznawczej.

W tym artykule pokazano, jak utworzyć narzędzia dla agenta sztucznej inteligencji do wyszukiwania danych nieustrukturyzowanych przy użyciu Frameworku Mosaic AI Agent. Nieustrukturyzowane metody pobierania umożliwiają agentom wykonywanie zapytań dotyczących źródeł danych bez struktury, takich jak corpus dokumentu, przy użyciu indeksów wyszukiwania wektorowego.

Aby dowiedzieć się więcej o narzędziach agentów, zobacz Stwórz narzędzia agenta AI.

Narzędzie do wyszukiwania wektorowego z funkcjami Catalog Unity

Poniższy przykład tworzy funkcję Catalog w silniku Unity dla narzędzia pobierania, które może zapytać o dane z indeksu Mosaic AI Vector Search.

Funkcja Catalog Unity databricks_docs_vector_search przeszukuje hipotetyczny indeks wyszukiwania wektorowego zawierający dokumentację Databricks. Opakowuje funkcję SQL Databricks vector_search() i używa aliasów page_content i metadata, aby dopasować jej wynik do retrieve'ra MLflow schema.

Notatka

Aby zapewnić zgodność z procesem pobierania schemaMLflow, wszystkie dodatkowe metadane columns należy dodać do metadatacolumn za pomocą funkcji mapowania SQL, a nie jako klucze wyjściowe najwyższego poziomu.

Uruchom następujący kod w notesie lub edytorze SQL.

CREATE OR REPLACE FUNCTION main.default.databricks_docs_vector_search (
  -- The agent uses this comment to determine how to generate the query string parameter.
  query STRING
  COMMENT 'The query string for searching Databricks documentation.'
) RETURNS TABLE
-- The agent uses this comment to determine when to call this tool. It describes the types of documents and information contained within the index.
COMMENT 'Executes a search on Databricks documentation to retrieve text documents most relevant to the input query.' RETURN
SELECT
  chunked_text as page_content,
  map('doc_uri', url, 'chunk_id', chunk_id) as metadata
FROM
  vector_search(
    -- Specify your Vector Search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

To narzędzie pobierania ma następujące zastrzeżenia:

• Narzędzie MLflow śledzi tę funkcję Unity Catalog jako zakres typu TOOL, a nie jako zakres typu RETRIEVER. W związku z tym aplikacje używające infrastruktury agentów, takie jak aplikacja do przeglądu agentów i AI Playground, nie będą wyświetlać szczegółów specyficznych dla retrievera, takich jak linki do fragmentów. Aby uzyskać więcej informacji na temat typów span, zobacz śledzenie MLflow Schema.
• Klienci SQL mogą limit maksymalną liczbę zwracanych wierszy lub bajtów. Aby zapobiec obcięciu danych, należy skrócić columnvalues zwracane przez funkcję UDF. Na przykład można użyć substring(chunked_text, 0, 8192), aby zmniejszyć rozmiar dużej zawartości columns i uniknąć obcinania wierszy podczas wykonywania.
• Ponieważ to narzędzie jest nakładką funkcji vector_search(), podlega tym samym ograniczeniom, co funkcja vector_search(). Zobacz Ograniczenia.

Jeśli ten przykład jest nieodpowiedni dla twojego przypadku użycia, zamiast tego utwórz narzędzie do wyszukiwania i pobierania wektorów przy użyciu niestandardowego kodu agenta.

Funkcja pobierania wyszukiwania wektorowego z kodem agenta (PyFunc)

Poniższy przykład tworzy wyszukiwanie wektorowe dla agenta w stylu PyFunc w kodzie agenta.

Ten przykład używa databricks-vectorsearch do utworzenia podstawowego modułu wyszukiwania, który wykonuje wyszukiwanie podobieństwa wektorowego z filtrami. Używa dekoratorów MLflow do włączenia śledzenia agentów.

Notatka

Aby zachować zgodność z pobieraczem MLflow schema, funkcja pobieracza powinna zwrócić typ dokumentu Document i użyć pola metadata w klasie Document, aby dodać dodatkowe atrybuty do zwróconego dokumentu, takie jak like doc_uri i similarity_score..

Użyj następującego kodu w module agenta lub notesie agenta.

import mlflow
import json

from mlflow.entities import Document
from typing import List, Dict, Any
from dataclasses import asdict
from databricks.vector_search.client import VectorSearchClient

class VectorSearchRetriever:
    """
    Class using Databricks Vector Search to retrieve relevant documents.
    """
    def __init__(self):
        self.vector_search_client = VectorSearchClient(disable_notice=True)
        # TODO: Replace this with the list of column names to return in the result when querying Vector Search
        self.columns = ["chunk_id", "text_column", "doc_uri"]
        self.vector_search_index = self.vector_search_client.get_index(
            index_name="catalog.schema.chunked_docs_index"
        )
        mlflow.models.set_retriever_schema(
            name="vector_search",
            primary_key="chunk_id",
            text_column="text_column",
            doc_uri="doc_uri"
        )

    @mlflow.trace(span_type="RETRIEVER", name="vector_search")
    def __call__(
        self,
        query: str,
        filters: Dict[Any, Any] = None,
        score_threshold = None
    ) -> List[Document]:
        """
        Performs vector search to retrieve relevant chunks.
        Args:
            query: Search query.
            filters: Optional filters to apply to the search. Filters must follow the Databricks Vector Search filter spec
            score_threshold: Score threshold to use for the query.

        Returns:
            List of retrieved Documents.
        """

        results = self.vector_search_index.similarity_search(
            query_text=query,
            columns=self.columns,
            filters=filters,
            num_results=5,
            query_type="ann"
        )

        documents = self.convert_vector_search_to_documents(
            results, score_threshold
        )
        return [asdict(doc) for doc in documents]

    @mlflow.trace(span_type="PARSER")
    def convert_vector_search_to_documents(
        self, vs_results, score_threshold
    ) -> List[Document]:

        docs = []
        column_names = [column["name"] for column in vs_results.get("manifest", {}).get("columns", [])]
        result_row_count = vs_results.get("result", {}).get("row_count", 0)

        if result_row_count > 0:
            for item in vs_results["result"]["data_array"]:
                metadata = {}
                score = item[-1]

                if score >= score_threshold:
                    metadata["similarity_score"] = score
                    for i, field in enumerate(item[:-1]):
                        metadata[column_names[i]] = field

                    page_content = metadata.pop("text_column", None)

                    if page_content:
                        doc = Document(
                            page_content=page_content,
                            metadata=metadata
                        )
                        docs.append(doc)

        return docs

Aby uruchomić program retriever, uruchom następujący kod w języku Python. Możesz opcjonalnie uwzględnić filtry wyszukiwania wektorowego w żądaniu, w celu filtrowania wyników.

retriever = VectorSearchRetriever()
query = "What is Databricks?"
filters={"text_column LIKE": "Databricks"},
results = retriever(query, filters=filters, score_threshold=0.1)

Set wykrywacz schema

Aby upewnić się, że retrievery są prawidłowo śledzone i poprawnie renderowane w aplikacjach podrzędnych, wywołaj mlflow.models.set_retriever_schema podczas definiowania agenta. Użyj set_retriever_schema, aby przypisać nazwy column w wynikach table do oczekiwanych pól w MLflow, takich jak primary_key, text_columni doc_uri.

# Define the retriever's schema by providing your column names
mlflow.models.set_retriever_schema(
    name="vector_search",
    primary_key="chunk_id",
    text_column="text_column",
    doc_uri="doc_uri"
    # other_columns=["column1", "column2"],
)

Możesz również określić dodatkowe columns w schema modułu pobierania, podając list nazw column z polem other_columns.

Jeśli masz wiele modułów pobierania, możesz zdefiniować wiele schematów przy użyciu unikatowych nazw dla każdego modułu pobierania schema.

Retriever schemaset podczas tworzenia agenta wpływa na aplikacje podrzędne i przepływy pracy, takie jak aplikacji przeglądu i zestawy oceny. W szczególności doc_uricolumn służy jako podstawowy identifier dla dokumentów zwracanych przez moduł retriever.

• Aplikacja recenzji wyświetla doc_uri, aby pomóc recenzentom ocenić odpowiedzi i śledzić pochodzenie dokumentów. Zobacz Przegląd interfejsu użytkownika aplikacji.
• Zestawy ewaluacyjne używane są doc_uri do porównywania wyników retrievera z wcześniej zdefiniowanymi zestawami danych ewaluacyjnych w celu określenia kompletności i dokładności retrievera. Zobacz zestawy oceny .

Śledź retrievera

Śledzenie MLflow zwiększa obserwowalność, przechwytując szczegółowe informacje na temat wykonywania agenta. Zapewnia ona sposób rejestrowania danych wejściowych, wyjściowych i metadanych skojarzonych z każdym pośrednim krokiem żądania, umożliwiając szybkie wskazanie źródła usterek i nieoczekiwanych zachowań.

W tym przykładzie użyto decoratora @mlflow.trace, aby utworzyć ślad dla modułu pobierania i analizatora. Aby uzyskać inne opcje konfigurowania metod śledzenia, zobacz śledzenie MLflow dla agentów.

Dekorator tworzy zakres , który zaczyna się, gdy funkcja jest wywoływana, a kończy przy zwróceniu wyniku funkcji. Narzędzie MLflow automatycznie rejestruje dane wejściowe i wyjściowe funkcji oraz wszelkie zgłoszone wyjątki.

Notatka

Użytkownicy biblioteki LangChain, LlamaIndex i OpenAI mogą używać automatycznego rejestrowania MLflow zamiast ręcznie definiować ślady za pomocą dekoratora. Zobacz Użyj autologowania, aby dodać ślady do swoich agentów.

...
@mlflow.trace(span_type="RETRIEVER", name="vector_search")
def __call__(self, query: str) -> List[Document]:
  ...

Aby upewnić się, że aplikacje końcowe, takie jak Ocena Agenta i Plac Zabaw AI, renderują ślad pobieracza prawidłowo, upewnij się, że dekorator spełnia następujące wymagania:

• Użyj span_type="RETRIEVER" i upewnij się, że funkcja zwraca obiekt List[Document]. Zobacz Retriever obejmuje.
• Nazwa śledzenia i nazwa retriever_schema muszą być zgodne, aby poprawnie skonfigurować śledzenie.

Filtrowanie wyników wyszukiwania wektorowego

Zakres wyszukiwania można limit do podzbioru danych przy użyciu filtru wyszukiwania wektorowego.

Parametr filters w VectorSearchRetriever definiuje warunki filtrowania przy użyciu specyfikacji filtru Databricks Vector Search.

filters = {"text_column LIKE": "Databricks"}

Wewnątrz metody __call__ słownik filtrów jest przekazywany bezpośrednio do funkcji similarity_search:

results = self.vector_search_index.similarity_search(
    query_text=query,
    columns=self.columns,
    filters=filters,
    num_results=5,
    query_type="ann"
)

Po początkowym filtrowaniu parametr score_threshold zapewnia dodatkowe filtrowanie, ustawiając minimalny wynik podobieństwa.

if score >= score_threshold:
    metadata["similarity_score"] = score

Końcowy wynik zawiera dokumenty spełniające warunki filters i score_threshold.

Następne kroki

Po utworzeniu narzędzia funkcji Unity Catalog dodaj narzędzie do agenta AI. Zobacz Dodaj narzędzia Catalog Unity do agentów.