Build und Tracing von Retriever-Tools für unstrukturierte Daten

Verwenden Sie das Mosaik AI Agent Framework, um Tools zu erstellen, mit denen KI-Agents unstrukturierte Daten wie eine Sammlung von Dokumenten abfragen können. Auf dieser Seite wird gezeigt, wie Sie:

• Lokales Entwickeln von Retrievern
• Erstellen eines Retrievers mithilfe von Unity-Katalogfunktionen
• Abfragen externer Vektorindizes
• Hinzufügen von MLflow-Ablaufverfolgung zur Beobachtbarkeit

Weitere Informationen zu Agenttools finden Sie unter KI-Agent-Tools.

Lokal Vektorsuch-Retriever-Tools mit der AI Bridge entwickeln

Die schnellste Möglichkeit, mit dem Erstellen eines Databricks Vector Search-Retriever-Tools zu beginnen, besteht darin, es lokal mithilfe von Databricks AI Bridge-Paketen zu entwickeln und zu testen, wie databricks-langchain und databricks-openai.

LangChain/LangGraph

Installieren Sie die neueste Version von databricks-langchain, die Databricks AI Bridge enthält.

%pip install --upgrade databricks-langchain

Der folgende Codeprototyp stellt ein Retriever-Tool bereit, das einen hypothetischen Vektorsuchindex abfragt und lokal an eine LLM bindet, damit Sie das Toolaufrufverhalten testen können.

Stellen Sie einen beschreibenden Text tool_description bereit, der dem Agent hilft, das Tool zu verstehen und zu bestimmen, wann es aufgerufen werden soll.

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-sonnet-4-5")
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?")

Für Szenarien, die entweder Direktzugriffsindizes oder Delta-Synchronisierungsindizes mit selbstverwalteten Einbettungen verwenden, müssen Sie das VectorSearchRetrieverTool benutzerdefinierte Einbettungsmodell und eine Textspalte konfigurieren und angeben. Weitere Informationen finden Sie unter Optionen zum Bereitstellen von Einbettungen.

Das folgende Beispiel zeigt, wie Sie ein VectorSearchRetrieverTool mit columns- und embedding-Schlüsseln konfigurieren.

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.
)

Weitere Details finden Sie in den API-Dokumenten für VectorSearchRetrieverTool.

OpenAI

Installieren Sie die neueste Version von databricks-openai, die Databricks AI Bridge enthält.

%pip install --upgrade databricks-openai

Der folgende Codeprototyp veranschaulicht einen Retriever, der einen hypothetischen Vektorsuchindex abfragt und in die GPT-Modelle von OpenAI integriert.

Stellen Sie einen beschreibenden Text tool_description bereit, der dem Agent hilft, das Tool zu verstehen und zu bestimmen, wann es aufgerufen werden soll.

Weitere Informationen zu OpenAI-Empfehlungen für Tools finden Sie in OpenAI Function Calling-Dokumentation.

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]
)

Für Szenarien, die entweder Direktzugriffsindizes oder Delta-Synchronisierungsindizes mit selbstverwalteten Einbettungen verwenden, müssen Sie das VectorSearchRetrieverTool benutzerdefinierte Einbettungsmodell und eine Textspalte konfigurieren und angeben. Weitere Informationen finden Sie unter Optionen zum Bereitstellen von Einbettungen.

Das folgende Beispiel zeigt, wie Sie ein VectorSearchRetrieverTool mit columns- und embedding-Schlüsseln konfigurieren.

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.
)

Weitere Details finden Sie in den API-Dokumenten für VectorSearchRetrieverTool.

Sobald Ihr lokales Tool bereit ist, können Sie es direkt als Teil Ihres Agentcodes produzieren oder zu einer Unity-Katalogfunktion migrieren, die eine bessere Auffindbarkeit und Governance bietet, aber bestimmte Einschränkungen hat.

Im folgenden Abschnitt wird gezeigt, wie Sie den Retriever zu einer Unity-Katalogfunktion migrieren.

Vektorsuch-Retriever-Tool mit Unity-Katalogfunktionen

Sie können eine Unity-Katalogfunktion erstellen, die eine Indexabfrage für die Mosaik-AI-Vektorsuche umschließt. Dieser Ansatz:

• Unterstützung für produktive Anwendungsfälle mit Governance und Auffindbarkeit
• Verwendet die SQL-Funktion vector_search() unter der Haube.
• Unterstützt automatisches MLflow-Tracing
  • Sie müssen die Ausgabe der Funktion mithilfe der Aliase und page_content an das metadata ausrichten.
  • Alle zusätzlichen Metadatenspalten müssen der Spalte mithilfe der metadataSQL-Zuordnungsfunktion und nicht als Ausgabeschlüssel der obersten Ebene hinzugefügt werden.

Führen Sie den folgenden Code in einem Notizbuch oder SQL-Editor aus, um die Funktion zu erstellen:

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
  )

Um dieses Retriever-Tool in Ihrem KI-Agent zu verwenden, schließen Sie es mit UCFunctionToolkitum. Dies ermöglicht automatisches Tracing durch MLflow, indem automatisch RETRIEVER Span-Typen in MLflow-Protokollen generiert werden.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Unity Catalog Retriever-Tools weisen die folgenden Vorbehalte auf:

• SQL-Clients können die maximale Anzahl zurückgegebener Zeilen oder Bytes einschränken. Um Datenabschneidung zu verhindern, sollten Sie die von der UDF zurückgegebenen Spaltenwerte kürzen. Sie können z. B. substring(chunked_text, 0, 8192) verwenden, um die Größe großer Inhaltsspalten zu verringern und zeilenabkürzungen während der Ausführung zu vermeiden.
• Da es sich bei diesem Tool um einen Wrapper für die vector_search()-Funktion handelt, unterliegt es den gleichen Einschränkungen wie die vector_search()-Funktion. Informationen finden Sie unter Einschränkungen.

Weitere Informationen zu UCFunctionToolkit finden Sie in Unity-Katalogdokumentation.

Retriever, der einen außerhalb von Databricks gehosteten Vektorindex abfragt

Wenn Ihr Vektorindex außerhalb von Azure Databricks gehostet wird, können Sie eine Unity-Katalogverbindung erstellen, um eine Verbindung mit dem externen Dienst herzustellen und die Verbindung in Ihrem Agentcode zu verwenden. Siehe Verbinden von KI-Agent-Tools mit externen Diensten.

Im folgenden Beispiel wird ein Retriever erstellt, der einen Vektorindex aufruft, welcher außerhalb von Databricks für einen PyFunc-gebundenen Agenten gehostet wird.

1. Erstellen Sie eine Unity-Katalogverbindung mit dem externen Dienst, in diesem Fall Azure.

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

2. Definieren Sie das Retriever-Tool im Agentcode mithilfe der Unity-Katalogverbindung. In diesem Beispiel werden MLflow-Decorator-Elemente verwendet, um die Agent-Nachverfolgung zu aktivieren.

   Hinweis

   Um dem MLflow-Retriever-Schema zu entsprechen, sollte die Retriever-Funktion ein List[Document]-Objekt zurückgeben und das metadata-Feld in der Document-Klasse verwenden, um dem zurückgegebenen Dokument zusätzliche Attribute hinzuzufügen, wie z. B. doc_uri und similarity_score. Siehe MLflow-Dokument.

   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. Führen Sie zum Ausführen des Retrievers den folgenden Python-Code aus. Sie können optional Vektorsuchfilter in die Anforderung einschließen, um Ergebnisse zu filtern.

   retriever = VectorSearchRetriever()
   query = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
   results = retriever(query, score_threshold=0.1)
   

Tracing zu einem Retriever hinzufügen

Fügen Sie die MLflow-Ablaufverfolgung hinzu, um den Retriever zu überwachen und zu debuggen. Mit der Ablaufverfolgung können Sie Eingaben, Ausgaben und Metadaten für jeden Schritt der Ausführung anzeigen.

Das vorherige Beispiel fügt den @mlflow.trace Dekorator sowohl zu den __call__ als auch zu den Parsing-Methoden hinzu. Der Dekorateur erstellt einen Bereich, der beginnt, wenn die Funktion aufgerufen wird, und endet, wenn sie zurückkehrt. MLflow zeichnet automatisch die Eingabe und Ausgabe der Funktion und alle ausgelösten Ausnahmen auf.

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]:
    ...

Um sicherzustellen, dass Downstreamanwendungen wie „Agent-Auswertung“ und „KI-Playground“ die Retriever-Nachverfolgung korrekt rendern, stellen Sie sicher, dass der Decorator die folgenden Anforderungen erfüllt:

• Verwenden Sie (https://mlflow.org/docs/latest/tracing/tracing-schema.html#retriever-spans) und stellen Sie sicher, dass die Funktion ein List[Document]-Objekt zurückgibt.
• Der Protokollname und der retriever_schema Name müssen übereinstimmen, um die Protokollierung korrekt zu konfigurieren. Im folgenden Abschnitt erfahren Sie, wie Sie das Retriever-Schema festlegen.

Festlegen des Retriever-Schemas, um die MLflow-Kompatibilität sicherzustellen

Wenn das vom Retriever oder span_type="RETRIEVER" zurückgegebene Tracing nicht dem MLflow Standard-Retriever-Schema entspricht, müssen Sie das zurückgegebene Schema manuell den von MLflow erwarteten Feldern zuordnen. Dadurch wird sichergestellt, dass MLflow Ihren Retriever ordnungsgemäß verfolgen und Darstellungen in nachgelagerten Anwendungen erzeugen kann.

So legen Sie das Retriever-Schema manuell fest:

1. Rufen Sie mlflow.models.set_retriever_schema auf, wenn Sie Ihren Agent definieren. Verwenden Sie set_retriever_schema, um die Spaltennamen in der zurückgegebenen Tabelle den erwarteten Feldern von MLflow wie primary_key, text_column und doc_uri zuzuordnen.

   # 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. Geben Sie zusätzliche Spalten im Schema des Retrievers an, indem Sie eine Liste der Spaltennamen mit dem other_columns Feld angeben.

3. Wenn Sie über mehrere Retriever verfügen, können Sie mehrere Schemas definieren, indem Sie eindeutige Namen für jedes Retriever-Schema verwenden.

Das während der Erstellung des Agenten festgelegte Retriever-Schema wirkt sich auf nachgeschaltete Anwendungen und Workflows aus, wie z. B. die Review-App und Evaluierungssätze. Insbesondere dient die spalte doc_uri als primärer Bezeichner für dokumente, die vom Retriever zurückgegeben werden.

• Die Überprüfungs-App zeigt die doc_uri an, damit Prüfer Antworten bewerten und Dokumentenherkunft nachverfolgen können. Weitere Informationen finden Sie unter Überprüfen der App-Benutzeroberfläche.
• Evaluierungs-Sets verwenden doc_uri, um die Retriever-Ergebnisse mit vordefinierten Evaluierungs-Datasets zu vergleichen, um den Recall und die Präzision des Retrievers zu bestimmen. Siehe Auswertungssätze (MLflow 2)

Nächste Schritte

Nachdem Sie ein Unity-Katalogtool erstellt haben, fügen Sie es Ihrem Agent hinzu. Siehe Erstellen eines Agent-Tools.