Freigeben über

Tools für unstrukturiertes Abrufen des KI-Agents

  • 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 Erstellen von KI-Agent-Tools.

Vektorsuche-Retriever-Tool mit Unity-Katalogfunktionen

Im folgenden Beispiel wird eine Unity-Katalogfunktion für ein Retriever-Tool erstellt, mit dem Daten aus einem Mosaik AI Vector Search-Indexabgefragt werden können.

Die Unity-Katalogfunktion databricks_docs_vector_search fragt einen hypothetischen Vektorsuchindex ab, der die Dokumentation zu Databricks enthält. Sie umschließt die SQL-Funktion vector_search() der Databricks-Funktion und verwendet die Aliase page_content und metadata, um die Ausgabe mit dem MLflow-Retrieverschemaabzugleichen.

Anmerkung

Um dem MLflow-Abrufschema zu entsprechen, müssen alle zusätzlichen Metadatenspalten der metadata-Spalte mithilfe der SQL-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.

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
  )

Dieses Retriever-Tool weist die folgenden Vorbehalte auf:

  • MLflow verfolgt diese Unity Catalog-Funktion als TOOL-Span-Typ und nicht als RETRIEVER-Span-Typ nach. Daher werden Agent-Framework-Downstreamanwendungen wie die Agent-Überprüfungs-App und KI Playground keine retrieverspezifischen Details wie Verknüpfungen zu Blöcken anzeigen. Weitere Informationen zu Span-Typen finden Sie unter MLflow-Ablaufverfolgungsschema.
  • SQL-Clients können die maximale Anzahl zurückgegebener Zeilen oder Bytes einschränken. Um das Abschneiden von Daten zu verhindern, sollten Sie Spaltenwerte abschneiden, die von UDF zurückgegeben werden. 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. Siehe Einschränkungen.

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

Vektor-Sucher-Retriever mit Agentencode (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.

Anmerkung

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.

Der Retriever-Schemasatz während der Agent-Erstellung wirkt sich auf Downstreamanwendungen und Workflows aus, z. B. auf die Prüf-App und auf Auswertungssä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 zum Einrichten von Ablaufverfolgungsmethoden finden Sie unter MLflow-Ablaufverfolgung für Agents.

Das Decorator-Element erstellt einen Span, der beginnt, wenn die Funktion aufgerufen wird, und endet, wenn sie zurückgegeben wird. MLflow zeichnet automatisch die Eingabe und Ausgabe der Funktion und alle ausgelösten Ausnahmen auf.

Anmerkung

LangChain-, LlamaIndex- und OpenAI-Bibliotheksbenutzende können die automatische Protokollierung von MLflow nutzen, anstatt die Ablaufverfolgungen manuell mit dem Decorator-Element 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.

Ergebnisse der Vektorsuche 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.