Erstellen eines Vektorsuch-Retriever-Tools

Erfahren Sie, wie Sie das Mosaik AI Agent Framework zum Erstellen von Retrievern verwenden. Ein Retriever ist ein Agenttooltyp, der relevante Dokumente mithilfe eines Vektorsuchindex findet und zurückgibt. Retriever sind eine Kernkomponente von RAG-Anwendungen (Retrieval Augmented Generation).

Anforderungen

• MLflow Document ist nur für MLflow Version 2.14.0 und höher verfügbar.
• Ein vorhandener Vektorsuchindex.

PyFunc-Retriever-Beispiel

Im folgenden Beispiel wird die Databricks-Vectorsearch verwendet, um einen einfachen Retriever zu erstellen, der eine Vektorsuche mit Filtern durchführt. Es verwendet MLflow-Dekoratoren, um die Agent-Ablaufverfolgung zu aktivieren.

Die Retriever-Funktion sollte einen Dokumenttyp zurückgeben. Verwenden Sie das metadata Feld in der Document-Klasse, um dem zurückgegebenen Dokument zusätzliche Attribute hinzuzufügen, z like doc_uri . B. 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. Optional können Sie 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

Um sicherzustellen, dass Retriever ordnungsgemäß nachverfolgt werden, rufen Sie mlflow.models.set_retriever_schema auf, wenn Sie Ihren Agent im Code definieren. Dient set_retriever_schema zum Zuordnen der Spaltennamen in der zurückgegebenen Tabelle zu den erwarteten Feldern von MLflow, z primary_key. B. , text_columnund 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"],
)

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

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

Nachverfolgen des Retrievers

Die MLflow-Ablaufverfolgung fügt Die Observability hinzu, indem detaillierte Informationen zur Ausführung Des Agents erfasst werden. Es bietet eine Möglichkeit, die Eingaben, Ausgaben und Metadaten aufzuzeichnen, die jedem Zwischenschritt einer Anforderung zugeordnet sind, sodass Sie die Quelle von Fehlern und unerwarteten Verhaltensweisen leicht anheften können.

In diesem Beispiel wird der @mlflow.trace-Dekoror verwendet, um eine Ablaufverfolgung für den Retriever und parser zu erstellen. Weitere Optionen zum Einrichten von Ablaufverfolgungsmethoden finden Sie unter MLflow-Ablaufverfolgung für Agents.

Der Dekorateur erstellt eine Spanne , die beginnt, wenn die Funktion aufgerufen wird und endet, wenn sie zurückgegeben wird. MLflow zeichnet automatisch die Eingabe und Ausgabe der Funktion und alle ausnahmen auf, die daraus ausgelöst wurden.

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

Um sicherzustellen, dass downstream-Anwendungen wie Agent Evaluation und AI Playground die Retriever-Ablaufverfolgung korrekt rendern, stellen Sie sicher, dass der Dekorierer die folgenden Anforderungen erfüllt:

• Verwenden Sie das Objekt, und stellen Sie span_type="RETRIEVER" sicher, dass die Funktion ein Objekt zurückgibt List[Document] . Siehe Retriever-Spanne.
• Der Ablaufverfolgungsname und der retriever_schema Name müssen übereinstimmen, um die Ablaufverfolgung ordnungsgemäß zu konfigurieren.

Filtervektorsuchergebnisse

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 Filterspezifikation "Databricks Vector Search".

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 stellt der score_threshold Parameter zusätzliche Filterung bereit, indem eine Mindestgleichheitsbewertung festgelegt wird.

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

Das Endergebnis enthält Dokumente, die die filters Bedingungen score_threshold erfüllen.

Retriever-Beispielanwendungen

Sehen Sie sich das Genai-Cookbook GitHub-Repository für KI-Agent-Beispiele an, die Retriever verwenden:

• LangChain-Retriever-Beispiel: Demo: Mosaik AI Agent Framework und Agent Evaluation
• PyFunc-Retriever-Beispiel: Agent-App-Beispielcode

Nächste Schritte

• Erstellen eines KI-Agents