Delen via

Een hulpprogramma voor het ophalen van vectorzoekopdrachten maken

  • Artikel

Belangrijk

Deze functie is beschikbaar als openbare preview.

Meer informatie over het gebruik van Mosaic AI Agent Framework om retrievers te maken. Een retriever is een type agenthulpprogramma waarmee relevante documenten worden gevonden en geretourneerd met behulp van een Vector Search-index. Retrievers zijn een kernonderdeel van RAG-toepassingen (Retrieval Augmented Generation).

Vereisten

  • MLflow is alleen beschikbaar op MLflow Document versie 2.14.0 en hoger.
  • Een bestaande Vector Search-index.

Voorbeeld van PyFunc-retriever

In het volgende voorbeeld wordt databricks-vectorsearch gebruikt om een eenvoudige retriever te maken waarmee een vectorzoekactie met filters wordt uitgevoerd. Er worden MLflow-decorators gebruikt om agenttracering in te schakelen.

De functie retriever moet een documenttype retourneren. Gebruik het metadata veld in de documentklasse om extra kenmerken toe te voegen aan het geretourneerde document, zoals like doc_uri en similarity_score.

Gebruik de volgende code in de agentmodule of agentnotebook.

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

Voer de volgende Python-code uit om de retriever uit te voeren. U kunt eventueel Vector Search-filters opnemen in de aanvraag om resultaten te filteren.

retriever = VectorSearchRetriever()
query = "What is Databricks?"
filters={"text_column LIKE": "Databricks"},
results = retriever(query, filters=filters, score_threshold=0.1)

Schema voor ophalen instellen

Als u ervoor wilt zorgen dat retrievers correct worden getraceerd, roept u mlflow.models.set_retriever_schema aan wanneer u uw agent in code definieert. Gebruik set_retriever_schema deze functie om de kolomnamen in de geretourneerde tabel toe te wijzen aan de verwachte velden van MLflow, zoals primary_key, text_columnen 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"],
)

Notitie

De doc_uri kolom is vooral belangrijk bij het evalueren van de prestaties van de retriever. doc_uri is de belangrijkste id voor documenten die door de retriever worden geretourneerd, zodat u deze kunt vergelijken met basis-waarheidsevaluatiesets. Zie Evaluatiesets.

U kunt ook extra kolommen opgeven in het schema van de retriever door een lijst met kolomnamen met het other_columns veld op te geven.

Als u meerdere retrievers hebt, kunt u meerdere schema's definiëren met behulp van unieke namen voor elk retriever-schema.

De retriever traceren

MLflow-tracering voegt waarneembaarheid toe door gedetailleerde informatie over de uitvoering van uw agent vast te leggen. Het biedt een manier om de invoer, uitvoer en metagegevens vast te leggen die zijn gekoppeld aan elke tussenliggende stap van een aanvraag, zodat u eenvoudig de bron van bugs en onverwacht gedrag kunt aanwijzen.

In dit voorbeeld wordt de @mlflow.trace decorator gebruikt om een tracering te maken voor de retriever en parser. Zie MLflow Tracing voor agents voor andere opties voor het instellen van traceringsmethoden.

De decorator maakt een periode die begint wanneer de functie wordt aangeroepen en eindigt wanneer deze wordt geretourneerd. MLflow registreert automatisch de invoer en uitvoer van de functie en eventuele uitzonderingen die daaruit zijn gegenereerd.

Notitie

Gebruikers van de LangChain-, LlamaIndex- en OpenAI-bibliotheek kunnen automatische aanmelding van MLflow gebruiken in plaats van traceringen handmatig te definiëren met de decorator. Zie Autologging gebruiken om traceringen toe te voegen aan uw agents.

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

Om ervoor te zorgen dat downstreamtoepassingen, zoals agentevaluatie en AI Playground, de tracering van de retriever correct weergeven, moet u ervoor zorgen dat de decorator voldoet aan de volgende vereisten:

  • Gebruik span_type="RETRIEVER" en zorg ervoor dat de functie object retourneert List[Document] . Zie Retriever spans.
  • De traceringsnaam en de retriever_schema naam moeten overeenkomen om de tracering correct te configureren.

Zoekresultaten filtervector

U kunt het zoekbereik beperken tot een subset van gegevens met behulp van een Vector Search-filter.

De filters parameter in VectorSearchRetriever definieert de filtervoorwaarden met behulp van de Databricks Vector Search-filterspecificatie.

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

In de __call__ methode wordt de filterswoordenlijst rechtstreeks doorgegeven aan de similarity_search functie:

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

Na het initiële filteren biedt de score_threshold parameter extra filters door een minimale gelijkenisscore in te stellen.

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

Het uiteindelijke resultaat bevat documenten die voldoen aan de filters en score_threshold voorwaarden.

Voorbeeldtoepassingen ophalen

Zie de GitHub-opslagplaats voor genai-cookbook voor voorbeelden van AI-agents die gebruikmaken van retrievers:

Volgende stappen