Narzędzia do tworzenia i śledzenia pobierania dla danych nieustrukturyzowanych

Użyj systemu Mosaic AI Agent Framework, aby budować narzędzia, które pozwalają agentom AI wyszukiwać dane niestrukturalne, takie jak zbiór dokumentów. Na tej stronie pokazano, jak:

• Lokalne opracowywanie mechanizmów pobierających
• Stwórz retriever przy użyciu funkcji Unity Catalog
• Wykonywanie zapytań względem indeksów wektorów zewnętrznych
• Dodaj trasowanie MLflow w celu zapewnienia obserwowalności

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

Lokalnie opracowywać systemy wyszukiwania wektorowego z AI Bridge

Najszybszym sposobem na rozpoczęcie budowania narzędzia do pobierania danych wektorowych w Databricks jest jego opracowanie i przetestowanie lokalnie przy użyciu pakietów Databricks AI Bridge, takich jak i databricks-langchain.

LangChain/LangGraph

Zainstaluj najnowszą wersję databricks-langchain, która zawiera Databricks AI Bridge.

%pip install --upgrade databricks-langchain

Poniższy kod tworzy prototyp narzędzia do pobierania, które wysyła zapytanie do hipotetycznego indeksu wyszukiwania wektorów i wiąże je lokalnie z modułem LLM, aby można było przetestować jego zachowanie wywołujące narzędzia.

Podaj opis, tool_description aby ułatwić agentowi zrozumienie narzędzia i określenie, kiedy go wywołać.

from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool.
vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation."
)

# Run a query against the vector search index locally for testing
vs_tool.invoke("Databricks Agent Framework?")

# Bind the retriever tool to your Langchain LLM of choice
llm = ChatDatabricks(endpoint="databricks-claude-3-7-sonnet")
llm_with_tools = llm.bind_tools([vs_tool])

# Chat with your LLM to test the tool calling functionality
llm_with_tools.invoke("Based on the Databricks documentation, what is Databricks Agent Framework?")

W przypadku scenariuszy korzystających z indeksów bezpośredniego dostępu lub indeksów usługi Delta Sync przy użyciu samodzielnie zarządzanych osadzeń należy skonfigurować VectorSearchRetrieverTool i określić niestandardowy model osadzenia oraz kolumnę tekstową. Zobacz opcje zapewniania embedowań.

W poniższym przykładzie pokazano, jak skonfigurować element VectorSearchRetrieverTool z kluczami columns oraz embedding.

from databricks_langchain import VectorSearchRetrieverTool
from databricks_langchain import DatabricksEmbeddings

embedding_model = DatabricksEmbeddings(
    endpoint="databricks-bge-large-en",
)

vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
  num_results=5, # Max number of documents to return
  columns=["primary_key", "text_column"], # List of columns to include in the search
  filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
  query_type="ANN", # Query type ("ANN" or "HYBRID").
  tool_name="name of the tool", # Used by the LLM to understand the purpose of the tool
  tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
  text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
  embedding=embedding_model # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Aby uzyskać dodatkowe informacje, zobacz dokumentację interfejsu API dla programu VectorSearchRetrieverTool.

OpenAI

Zainstaluj najnowszą wersję databricks-openai, która zawiera Databricks AI Bridge.

%pip install --upgrade databricks-openai

Poniższy kod tworzy prototypy programu retriever, który wysyła zapytanie do hipotetycznego indeksu wyszukiwania wektorów i integruje go z modelami GPT platformy OpenAI.

Podaj opis, tool_description aby ułatwić agentowi zrozumienie narzędzia i określenie, kiedy go wywołać.

Aby uzyskać więcej informacji na temat zaleceń OpenAI dotyczących narzędzi, zobacz dokumentację Dotyczącą Wywoływania Funkcji OpenAI.

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

# Initialize OpenAI client
client = OpenAI(api_key=<your_API_key>)

# Initialize the retriever tool
dbvs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation"
)

messages = [
  {"role": "system", "content": "You are a helpful assistant."},
  {
    "role": "user",
    "content": "Using the Databricks documentation, answer what is Spark?"
  }
]
first_response = client.chat.completions.create(
  model="gpt-4o",
  messages=messages,
  tools=[dbvs_tool.tool]
)

# Execute function code and parse the model's response and handle function calls.
tool_call = first_response.choices[0].message.tool_calls[0]
args = json.loads(tool_call.function.arguments)
result = dbvs_tool.execute(query=args["query"])  # For self-managed embeddings, optionally pass in openai_client=client

# Supply model with results – so it can incorporate them into its final response.
messages.append(first_response.choices[0].message)
messages.append({
  "role": "tool",
  "tool_call_id": tool_call.id,
  "content": json.dumps(result)
})
second_response = client.chat.completions.create(
  model="gpt-4o",
  messages=messages,
  tools=[dbvs_tool.tool]
)

W przypadku scenariuszy korzystających z indeksów bezpośredniego dostępu lub indeksów usługi Delta Sync przy użyciu samodzielnie zarządzanych osadzeń należy skonfigurować VectorSearchRetrieverTool i określić niestandardowy model osadzenia oraz kolumnę tekstową. Zobacz opcje zapewniania embedowań.

W poniższym przykładzie pokazano, jak skonfigurować element VectorSearchRetrieverTool z kluczami columns oraz embedding.

from databricks_openai import VectorSearchRetrieverTool

vs_tool = VectorSearchRetrieverTool(
    index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
    num_results=5, # Max number of documents to return
    columns=["primary_key", "text_column"], # List of columns to include in the search
    filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
    query_type="ANN", # Query type ("ANN" or "HYBRID").
    tool_name="name of the tool", # Used by the LLM to understand the purpose of the tool
    tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
    text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
    embedding_model_name="databricks-bge-large-en" # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Aby uzyskać dodatkowe informacje, zobacz dokumentację interfejsu API dla programu VectorSearchRetrieverTool.

Gdy narzędzie lokalne będzie gotowe, możesz bezpośrednio wdrożyć go do środowiska produkcyjnego w ramach kodu agenta lub przeprowadzić migrację do funkcji katalogu Unity, która zapewnia lepszą widoczność i zarządzanie, ale ma pewne ograniczenia.

W sekcji poniżej pokazano, jak przeprowadzić migrację programu Retriever do funkcji katalogu Unity.

narzędzie wyszukiwania wektorów z funkcjami katalogu Unity

Możesz utworzyć funkcję Unity Catalog, która otacza zapytanie indeksu wektora Mosaic AI. Takie podejście:

• Wspiera przypadki użycia w środowisku produkcyjnym z zarządzaniem i odkrywalnością.
• Używa funkcji sql vector_search() pod maską
• Obsługuje automatyczne śledzenie MLflow
  • Dane wyjściowe funkcji należy wyrównać do schematu modułu pobierania MLflow przy użyciu aliasów page_content i metadata.
  • Wszelkie dodatkowe kolumny metadanych należy dodać do metadata kolumny przy użyciu funkcji mapy SQL, a nie jako klucze wyjściowe najwyższego poziomu.

Uruchom następujący kod w notesie lub edytorze SQL, aby utworzyć funkcję:

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
  )

Aby użyć tego narzędzia pobierania w agencie sztucznej inteligencji, owiń je w UCFunctionToolkit. Umożliwia to automatyczne śledzenie za pośrednictwem platformy MLflow poprzez automatyczne generowanie typów zakresów w dziennikach MLflow.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

toolkit = UCFunctionToolkit(
    function_names=[
        "main.default.databricks_docs_vector_search"
    ]
)
tools = toolkit.tools

Narzędzia pobierania wykazu katalogu Unity mają następujące uwagi:

• Klienci SQL mogą ograniczyć maksymalną liczbę zwracanych wierszy lub bajtów. Aby zapobiec obcięciu danych, należy przyciąć wartości kolumn zwracane przez UDF. Na przykład można użyć substring(chunked_text, 0, 8192), aby zmniejszyć rozmiar dużych kolumn zawartości 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.

Aby uzyskać więcej informacji na temat UCFunctionToolkit, zobacz dokumentację Unity Catalog .

Usługa Retriever, która wykonuje zapytania dotyczące indeksu wektorowego hostowanego poza programem Databricks

Jeśli indeks wektorowy jest hostowany poza usługą Azure Databricks, możesz utworzyć połączenie Unity Catalog, aby nawiązać połączenie z usługą zewnętrzną i wykorzystać je w swoim kodzie agenta. Zobacz Połącz narzędzia agenta AI z usługami zewnętrznymi.

W poniższym przykładzie tworzony jest retriever, który wywołuje indeks wektorowy hostowany poza usługą Databricks dla agenta w wersji PyFunc.

1. Utwórz połączenie katalogu Unity do usługi zewnętrznej, w tym przypadku Azure.

   CREATE CONNECTION ${connection_name}
   TYPE HTTP
   OPTIONS (
     host 'https://example.search.windows.net',
     base_path '/',
     bearer_token secret ('<secret-scope>','<secret-key>')
   );
   

2. Zdefiniuj narzędzie wyszukiwania w kodzie agenta przy użyciu połączenia katalogu Unity. W tym przykładzie użyto dekoratorów MLflow do włączenia śledzenia działań agentów.

   Notatka

   Aby zachować zgodność ze schematem pobierania MLflow, funkcja pobierająca powinna zwrócić List[Document] obiekt i użyć metadata pola w klasie Document, aby dodać atrybuty do zwróconego dokumentu, na przykład doc_uri i similarity_score. Zobacz dokument MLflow.

   import mlflow
   import json
   
   from mlflow.entities import Document
   from typing import List, Dict, Any
   from dataclasses import asdict
   
   class VectorSearchRetriever:
     """
     Class using Databricks Vector Search to retrieve relevant documents.
     """
   
     def __init__(self):
       self.azure_search_index = "hotels_vector_index"
   
     @mlflow.trace(span_type="RETRIEVER", name="vector_search")
     def __call__(self, query_vector: List[Any], score_threshold=None) -> List[Document]:
       """
       Performs vector search to retrieve relevant chunks.
       Args:
         query: Search query.
         score_threshold: Score threshold to use for the query.
   
       Returns:
         List of retrieved Documents.
       """
       from databricks.sdk import WorkspaceClient
       from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
   
       json = {
         "count": true,
         "select": "HotelId, HotelName, Description, Category",
         "vectorQueries": [
           {
             "vector": query_vector,
             "k": 7,
             "fields": "DescriptionVector",
             "kind": "vector",
             "exhaustive": true,
           }
         ],
       }
   
       response = (
         WorkspaceClient()
         .serving_endpoints.http_request(
           conn=connection_name,
           method=ExternalFunctionRequestHttpMethod.POST,
           path=f"indexes/{self.azure_search_index}/docs/search?api-version=2023-07-01-Preview",
           json=json,
         )
         .text
       )
   
       documents = self.convert_vector_search_to_documents(response, 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 = []
   
       for item in vs_results.get("value", []):
         score = item.get("@search.score", 0)
   
         if score >= score_threshold:
           metadata = {
             "score": score,
             "HotelName": item.get("HotelName"),
             "Category": item.get("Category"),
           }
   
           doc = Document(
             page_content=item.get("Description", ""),
             metadata=metadata,
             id=item.get("HotelId"),
           )
           docs.append(doc)
   
       return docs
   

3. 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 = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
   results = retriever(query, score_threshold=0.1)
   

Dodawanie śledzenia do elementu wyszukiwania

Dodaj śledzenie MLflow, aby monitorować i debugować swojego retrievera. Śledzenie umożliwia wyświetlanie danych wejściowych, wyjściowych i metadanych dla każdego kroku wykonywania.

W poprzednim przykładzie dodano dekorator @mlflow.trace zarówno do __call__, jak i metod analizy. Dekorator tworzy przedział, który rozpoczyna się w momencie wywołania funkcji i kończy się po jej zakończeniu. Narzędzie MLflow automatycznie rejestruje dane wejściowe i wyjściowe funkcji oraz wszelkie zgłoszone wyjątki.

Notatka

Użytkownicy bibliotek LangChain, LlamaIndex i OpenAI mogą używać automatycznego rejestrowania MLflow, a także ręcznie definiować ślady przy użyciu dekoratora. Zobacz Instrumentacja aplikacji: metody śledzenia.

import mlflow
from mlflow.entities import Document

## This code snippet has been truncated for brevity, see the full retriever example above
class VectorSearchRetriever:
  ...

  # Create a RETRIEVER span. The span name must match the retriever schema name.
  @mlflow.trace(span_type="RETRIEVER", name="vector_search")
  def __call__(...) -> List[Document]:
    ...

  # Create a PARSER span.
  @mlflow.trace(span_type="PARSER")
  def parse_results(...) -> List[Document]:
    ...

Aby upewnić się, że aplikacje końcowe, takie jak Ocena Agenta i AI Playground, poprawnie renderują ślad pobieracza, należy zadbać o to, aby dekorator spełniał następujące wymagania:

• Użyj (https://mlflow.org/docs/latest/tracing/tracing-schema.html#retriever-spans) i upewnij się, że funkcja zwraca obiekt List[Document].
• Nazwa śledzenia i retriever_schema nazwa muszą być zgodne, aby poprawnie skonfigurować ślad. Zapoznaj się z poniższą sekcją, aby dowiedzieć się, jak ustawić schemat programu retriever.

Ustaw schemat wyszukiwania, aby zapewnić zgodność z MLflow

Jeśli ślad zwrócony z modułu pobierania lub span_type="RETRIEVER" nie jest zgodny ze standardowym schematem modułu pobierania MLflow, należy ręcznie zamapować zwrócony schemat na oczekiwane pola MLflow. Dzięki temu narzędzie MLflow może prawidłowo śledzić dane pobierania i renderować ślady w aplikacjach podrzędnych.

Aby ręcznie ustawić schemat programu retriever:

1. Wywołaj mlflow.models.set_retriever_schema podczas definiowania agenta. Użyj set_retriever_schema, aby zmapować nazwy kolumn w zwróconej tabeli na oczekiwane pola MLflow, takie jak primary_key, text_column i 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"],
   )
   

2. Określ dodatkowe kolumny w schemacie modułu pobierania, podając listę nazw kolumn z polem other_columns .

3. Jeśli masz wiele retrieverów, możesz zdefiniować wiele schematów, używając unikatowych nazw dla każdego schematu retrievera.

Schemat pobierania ustawiony podczas tworzenia agenta wpływa na aplikacje zależne i przepływy pracy, takie jak aplikacja przeglądowa i zestawy ocen. W szczególności kolumna doc_uri służy jako podstawowy identyfikator dokumentów zwracanych przez program retriever.

• Apka recenzji wyświetla doc_uri, aby pomóc recenzentom ocenić odpowiedzi i śledzić pochodzenie dokumentów. Zobacz Przegląd interfejsu użytkownika aplikacji.
• Zestawy oceny służą doc_uri do porównywania wyników modułu pobierającego ze wstępnie zdefiniowanymi zestawami danych oceny w celu określenia jego czułości i precyzji. Zobacz Zestawy ewaluacyjne (MLflow 2).

Następne kroki

Po utworzeniu retrievera, ostatnim krokiem jest integracja z definicją agenta sztucznej inteligencji. Dowiedz się, jak dodać narzędzie do agenta, zobacz Dodawanie narzędzi Unity Catalog do agentów.