Freigeben über

Unstrukturierte KI-Agentenwerkzeuge für das Datenabrufen

  • Artikel

Wichtig

Dieses Feature befindet sich in der Public Preview.

In diesem Artikel wird gezeigt, wie Sie KI-Agent-Tools für den unstrukturierten Datenabruf mithilfe des Mosaik AI Agent Frameworks erstellen. Unstrukturierte Retriever ermöglichen Agents das Abfragen unstrukturierter Datenquellen, z. B. eines Dokumentkorpus, mithilfe von Vektorsuchindizes.

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

Lokale Entwicklung von Abrufertools für die Vektorsuche mit AI Bridge

Die einfachste Möglichkeit, mit der Entwicklung eines Databricks Vector Search-Retriever-Tools zu beginnen, ist lokal. Verwenden Sie Databricks AI Bridge-Pakete wie databricks-langchain und databricks-openai, um einem Agent schnell Abruffunktionen hinzuzufügen und mit Abfrageparametern zu experimentieren. Dieser Ansatz ermöglicht eine schnelle Iteration während der ersten Entwicklung.

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. Weitere Informationen finden Sie unterAbrufertool für die Vektorsuche mit Unity Catalog-Funktionen.

LangChain/LangGraph

Der folgende Codeprototyp stellt ein Retriever-Tool her und bindet es lokal an ein LLM, damit Sie mit dem Agent chatten können, um sein Toolaufrufverhalten zu testen.

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

%pip install --upgrade databricks-langchain

Hinweis

Beim Initialisieren des VectorSearchRetrieverToolsind die Argumente text_column und embedding für Delta-Synchronisierungsindizes mit selbstverwalteten Einbettungen und Direct Vector Access-Indizes erforderlich. Weitere Informationen finden Sie unter Optionen zum Bereitstellen von Einbettungen.

from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool
vs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")

# 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-meta-llama-3-1-70b-instruct")
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?")

Übergeben Sie zum Anpassen von Toolaufrufen zusätzliche Argumente an VectorSearchRetrieverTool:

from databricks_langchain import VectorSearchRetrieverTool

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

OpenAI

Der folgende Codeprototyp veranschaulicht ein Vektorsuch-Retriever-Tool und integriert es in die GPT-Modelle von OpenAI.

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

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

%pip install --upgrade databricks-openai

Hinweis

Beim Initialisieren des VectorSearchRetrieverToolsind die Argumente text_column und embedding für Delta-Synchronisierungsindizes mit selbstverwalteten Einbettungen und Direct Vector Access-Indizes erforderlich. Weitere Informationen finden Sie unter Optionen zum Bereitstellen von Einbettungen.

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

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

# Call model with VectorSearchRetrieverTool defined
dbvs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")
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]
)

Übergeben Sie zum Anpassen von Toolaufrufen zusätzliche Argumente an VectorSearchRetrieverTool:

from databricks_openai import VectorSearchRetrieverTool

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

Vektorsuch-Retriever-Tool mit Unity-Katalogfunktionen

Im folgenden Beispiel wird ein Retriever-Tool mit einer Unity Catalog-Funktion erstellt, um Daten aus einem Mosaik AI Vector Search Indexabzufragen.

Die Unity-Katalogfunktion databricks_docs_vector_search fragt einen hypothetischen Vektorsuchindex ab, der die Dokumentation zu Databricks enthält. Diese Funktion umschließt die SQL-Funktion databricks vector_search() und richtet die Ausgabe an dem MLflow-Retrieverschemaaus. mit den Aliasen page_content und metadata.

Hinweis

Um dem MLflow-Retrieverschema zu entsprechen, müssen alle zusätzlichen Metadatenspalten mithilfe der SQL-Zuordnungsfunktionzur metadata-Spalte und nicht als übergeordnete Ausgabeschlüssel 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. Dadurch wird die automatische Ablaufverfolgung über MLflow ermöglicht.

Die MLflow-Ablaufverfolgung erfasst detaillierte Ausführungsinformationen für Anwendungen mit generativer KI. Es protokolliert Eingaben, Ausgaben und Metadaten für jeden Schritt, um Probleme zu debuggen und die Leistung zu analysieren.

Wenn UCFunctionToolkit verwendet wird, generieren Abrufer automatisch RETRIEVER-Spannentypen in MLflow-Protokollen, wenn ihre Ausgabe dem MLflow-Abruferschema entspricht. Weitere Informationen finden Sie unter MLflow-Ablaufverfolgungsschema.

Weitere Informationen zu UCFunctionToolkit finden Sie in Unity-Katalogdokumentation.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Dieses Retriever-Tool weist 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. Weitere Informationen finden Sie unter Einschränkungen.

Wenn dieses Beispiel für Ihren Anwendungsfall nicht geeignet ist, erstellen Sie stattdessen ein Vektorsuch-Retriever-Tool mit benutzerdefiniertem Agent-Code.

Abrufer für die Vektorsuche mit Agentcode (PyFunc)

Im folgenden Beispiel wird ein Vektorsuche-Retriever für einen PyFunc-Agent im Code eines Agents erstellt.

In diesem Beispiel wird databricks-vectorsearch- verwendet, um einen einfachen Retriever zu erstellen, der eine Vector Search-Ähnlichkeitssuche mit Filterndurchführt. Es werden MLflow-Decorator-Elemente verwendet, um die Agent-Nachverfolgung zu aktivieren.

Hinweis

Um dem MLflow-Retriever-Schema zu entsprechen, sollte die Retriever-Funktion einen Document Typ zurückgeben und das feld metadata in der Document-Klasse verwenden, um dem zurückgegebenen Dokument zusätzliche Attribute hinzuzufügen, z. B. like doc_uri und similarity_score.

Verwenden Sie den folgenden Code im Agentmodul oder Agent-Notizbuch.

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

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 = "What is Databricks?"
filters={"text_column LIKE": "Databricks"},
results = retriever(query, filters=filters, score_threshold=0.1)

Festlegen des Retriever-Schemas

Rufen Sie mlflow.models.set_retriever_schema auf, wenn Sie Ihren Agent definieren, um sicherzustellen, dass Retriever ordnungsgemäß nachverfolgt und in Downstreamanwendungen reibungslos gerendert werden. Verwenden Sie set_retriever_schema, um die Spaltennamen in der zurückgegebenen Tabelle den erwarteten Feldern von MLflow wie primary_key, text_columnund doc_urizuzuordnen.

# 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"],
)

Sie können auch zusätzliche Spalten im Schema des Retrievers angeben, indem Sie eine Liste der Spaltennamen mit dem feld other_columns angeben.

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.
  • Auswertungssätze verwenden doc_uri, um Retriever-Ergebnisse mit vordefinierten Auswertungsdatasets zu vergleichen, um den Abruf und die Genauigkeit des Retrievers zu bestimmen. Weitere Informationen finden Sie unter Auswertungssätze.

Nachverfolgen des Retrievers

Die MLflow-Ablaufverfolgung fügt Einblicke hinzu, indem die Ausführung Ihres Agents detailliert erfasst wird. Sie bietet eine Möglichkeit zum Aufzeichnen der Eingaben, Ausgaben und Metadaten, die jedem Zwischenschritt einer Anforderung zugeordnet sind, sodass Sie schnell die Quelle von Fehlern und unerwarteten Verhaltensweisen anheften können.

In diesem Beispiel wird dem @mlflow.trace-Decorator-Element verwendet, um eine Ablaufverfolgung für den Retriever und den Parser zu erstellen. Weitere Optionen zur Einrichtung von Tracking-Methoden finden Sie unter MLflow-Tracking für Agenten.

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.

Hinweis

LangChain-, LlamaIndex- und OpenAI-Bibliotheksbenutzer können die automatische MLflow-Protokollierung verwenden, anstatt Abläufe mit dem Decorator manuell zu definieren. Weitere Informationen finden Sie unter Verwenden der automatischen Protokollierung zum Hinzufügen von Ablaufverfolgungen zu Ihren Agents.

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

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

  • Verwenden Sie span_type="RETRIEVER", und stellen Sie sicher, dass die Funktion List[Document] Objekt zurückgibt. Weitere Informationen finden Sie unter Retriever-Span.
  • Der Ablaufverfolgungsname und der retriever_schema-Name müssen übereinstimmen, um die Ablaufverfolgung korrekt zu konfigurieren.

Suchergebnisse nach Vektoren filtern

Sie können den Suchbereich auf eine Teilmenge von Daten beschränken, indem Sie einen Vektorsuchfilter verwenden.

Der filters Parameter in VectorSearchRetriever definiert die Filterbedingungen mithilfe der Databricks Vector Search Filterspezifikation.

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

Innerhalb der __call__-Methode wird das Filterwörterbuch direkt an die similarity_search-Funktion übergeben:

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

Nach der anfänglichen Filterung bietet der score_threshold-Parameter zusätzliche Filterung, indem eine Mindestgleichheitsbewertung festgelegt wird.

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

Das Endergebnis enthält Dokumente, die den filters und score_threshold Bedingungen entsprechen.

Nächste Schritte

Nachdem Sie ein Unity Catalog-Funktions-Agent-Tool erstellt haben, fügen Sie das Tool einem KI-Agent hinzu. Weitere Informationen finden Sie unter Hinzufügen von Unity Catalog-Tools zu Agents.