Dela via

Ai-agentverktyg för ostrukturerad hämtning

  • Artikel

Viktig

Den här funktionen finns i offentlig förhandsversion.

Den här artikeln visar hur du skapar AI-agentverktyg för ostrukturerad datahämtning med hjälp av Mosaic AI Agent Framework. Ostrukturerade sökmotorer gör det möjligt för agenter att fråga ostrukturerade datakällor, till exempel en dokumentkorpus, genom att använda vektorsökningsindex.

Mer information om agentverktyg finns i Skapa AI-agentverktyg.

Vektorsökningsverktyget med Unity Catalog-funktioner

I följande exempel skapas en Unity Catalog-funktion för ett upphämtningsverktyg som kan ställa frågor mot data från ett Mosaic AI Vector Search-index.

Funktionen Unity Catalogdatabricks_docs_vector_search frågar ett hypotetiskt Vector Search-index som innehåller Dokumentation om Databricks. Den omsluter funktionen Databricks SQL vector_search() och använder aliasen page_content och metadata för att matcha dess utdata med MLflow retriever schema.

Not

Om du vill följa MLflow Retriever-schemamåste eventuella ytterligare metadata columns läggas till i metadatacolumn med hjälp av SQL-kartfunktioneni stället för som utdatanycklar på den översta nivån.

Kör följande kod i en notebook- eller SQL-redigerare.

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
  )

Det här verktyget för hämtning har följande varningar:

  • MLflow spårar den här Unity-Catalog-funktionen som en TOOL span-typ i stället för en RETRIEVER span-typ. Därför visar underordnade Agent Framework-program som agentgranskningsappen och AI Playground inte information som är specifik för hämtning, till exempel länkar till segment. Mer information om intervalltyper finns i MLflow Tracing Schema.
  • SQL-klienter kan limit det maximala antalet rader eller byte som returneras. För att förhindra datatrunkering bör du trunkera columnvalues som returneras av UDF. Du kan till exempel använda substring(chunked_text, 0, 8192) för att minska storleken på omfattande innehåll columns och undvika avkortning av rader under körningen.
  • Eftersom det här verktyget är en omslutning för funktionen vector_search() omfattas det av samma begränsningar som funktionen vector_search(). Se begränsningar.

Om det här exemplet är olämpligt för ditt användningsfall skapar du ett verktyg för vektorsökningshämtning med hjälp av anpassad agentkod i stället.

Vektorsökningshämtare med agentkod (PyFunc)

I följande exempel skapas en vektorsökhämtare för en PyFunc-anpassad agent i agentkoden.

I det här exemplet används databricks-vectorsearch för att skapa en grundläggande retriever som utför en likhetssökning i Vector Search med filter. Den använder MLflow-dekoratörer för att aktivera agentspårning.

Not

Om du vill följa MLflow Retriever-schemaska funktionen retriever returnera en dokumenttyp och använda fältet metadata i klassen Dokument för att lägga till ytterligare attribut i det returnerade dokumentet, till exempel like doc_uri och similarity_score.

Använd följande kod i agentmodulen eller agentanteckningsboken.

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

Kör följande Python-kod för att köra hämtaren. Du kan också inkludera filter för vektorsökning i begäran om att filtrera resultat.

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

Set retriever schema

För att säkerställa att hämtningarna spåras korrekt och återges korrekt i underordnade program anropar du mlflow.models.set_retriever_schema när du definierar din agent. Använd set_retriever_schema för att mappa de column namnen i den returnerade table till MLflows förväntade fält, till exempel primary_key, text_columnoch 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"],
)

Du kan också ange ytterligare columns i hämtarens schema genom att ange en list med column namn med fältet other_columns.

Om du har flera sökhundar, kan du definiera flera scheman med hjälp av unika namn för varje sökhund schema.

Retrievern schemaset vid skapandet av agenten påverkar efterföljande applikationer och arbetsflöden, till exempel granskningsappen och utvärderingsuppsättningar. Mer specifikt fungerar doc_uricolumn som den primära identifier för dokument som returneras av hämtaren.

  • granskningsappen visar doc_uri för att hjälpa granskare att utvärdera svar och spåra dokument ursprung. Se Granska appgränssnittet.
  • Utvärderingsuppsättningar används doc_uri för att jämföra resultaten från sökmotorn med fördefinierade utvärderingsdatauppsättningar för att fastställa sökmotorns träffsäkerhet och precision. Se Utvärderingsuppsättningar.

Spåra hämtningen

MLflow-spårning lägger till observerbarhet genom att samla in detaljerad information om agentens körning. Det ger ett sätt att registrera indata, utdata och metadata som är associerade med varje mellanliggande steg i en begäran, så att du snabbt kan hitta källan till buggar och oväntade beteenden.

I det här exemplet används @mlflow.trace-dekoratören för att skapa en spårning för sökare och parser. Andra alternativ för att konfigurera spårningsmetoder finns i MLflow Tracing för agenter.

Dekoratören skapar ett span som startar när funktionen anropas och slutar när den returneras. MLflow registrerar automatiskt funktionens indata och utdata och eventuella undantag som genereras.

Not

Användare av LangChain-, LlamaIndex- och OpenAI-bibliotek kan använda automatisk MLflow-loggning i stället för att manuellt definiera spårningar med dekoratören. Se Använd automatisk loggning för att lägga till spårningar till dina agenter.

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

Säkerställ att dekoratören uppfyller följande krav för att säkerställa att nedströmsapplikationer som Agent Evaluation och AI Playground återger upphämtarens spårning korrekt.

  • Använd span_type="RETRIEVER" och se till att funktionen returnerar List[Document] objekt. Se Retriever sträcker sig över.
  • Spårningsnamnet och retriever_schema namn måste matcha för att kunna konfigurera spårningen korrekt.

Sökresultat för filtervektor

Du kan limit sökomfånget till en delmängd data med hjälp av ett vektorsökningsfilter.

Parametern filters i VectorSearchRetriever definierar filtervillkoren med filterspecifikationen Databricks Vector Search.

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

I metoden __call__ skickas filterordlistan direkt till funktionen similarity_search:

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

Efter inledande filtrering ger parametern score_threshold ytterligare filtrering genom att ange en minsta likhetspoäng.

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

Slutresultatet innehåller dokument som uppfyller villkoren för filters och score_threshold.

Nästa steg

När du har skapat ett Unity Catalog-funktionsagentverktyg lägger du till verktyget i en AI-agent. Se Lägg till Unity Catalog-verktyg till agenter.