Ai-agentverktyg för ostrukturerad hämtning

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 AI-agentverktyg.

Utveckla verktyg för vektorsökning lokalt med AI Bridge

Det enklaste sättet att börja utveckla ett Databricks Vector Search retriever-verktyg är lokalt. Använd Databricks AI Bridge-paket som databricks-langchain och databricks-openai för att snabbt lägga till hämtningsfunktioner i en agent och experimentera med frågeparametrar. Den här metoden möjliggör snabb iteration under den inledande utvecklingen.

När det lokala verktyget är klart kan du direkt produktionsanpassa det som en del av din agentkod eller migrera det till en Unity Catalog-funktion, vilket ger bättre identifiering och styrning men har vissa begränsningar. Se verktyget Vector Search retriever medUnity Catalog-funktioner.

LangChain/LangGraph

Följande kod konstruerar en prototyp av ett hämtningsverktyg och binder det lokalt till en LLM så att du kan interagera med agenten för att testa hur det använder verktyget.

Installera den senaste versionen av databricks-langchain som innehåller Databricks AI Bridge.

%pip install --upgrade databricks-langchain

Not

När du initierar VectorSearchRetrieverToolkrävs argumenten text_column och embedding för Delta Sync-index med självhanterade inbäddningar och Direktvektoråtkomstindex. Se alternativ för att tillhandahålla inbäddningar.

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?")

Om du vill anpassa verktygsanrop skickar du ytterligare argument till 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

Följande kod prototyper ett vektorsökverktyg och integrerar det med OpenAI:s GPT-modeller.

Mer information om OpenAI-rekommendationer för verktyg finns i dokumentation om OpenAI-funktionsanrop.

Installera den senaste versionen av databricks-openai som innehåller Databricks AI Bridge.

%pip install --upgrade databricks-openai

Not

När du initierar VectorSearchRetrieverToolkrävs argumenten text_column och embedding för Delta Sync-index med självadministrerade inbäddningar och direkta vektoråtkomstindex. Se alternativ för att tillhandahålla inbäddningar.

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]
)

Om du vill anpassa verktygsanrop skickar du ytterligare argument till 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.
)

Retriever-verktyget för vektorsökning med Unity Catalog-funktioner

I följande exempel skapas verktyget retriever med hjälp av en Unity Catalog-funktion för att fråga efter data från ett Mosaic AI Vector Search-index.

Funktionen Unity Catalog databricks_docs_vector_search frågar ett hypotetiskt Vector Search-index som innehåller Dokumentation om Databricks. Den här funktionen omsluter Databricks SQL-funktionen vector_search() och justerar dess utmatning med MLflow retrieverschema. med hjälp av aliasen page_content och metadata.

Not

Om du vill följa MLflow Retriever-schemat måste eventuella ytterligare metadatakolumner läggas till i kolumnen metadata med hjälp av SQL-kartfunktioni stället för som utdatanycklar på den översta nivån.

Kör följande kod i en notebook- eller SQL-redigerare för att skapa funktionen:

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
  )

Om du vill använda det här hämtningsverktyget i AI-agenten omsluter du det med UCFunctionToolkit. Detta möjliggör automatisk spårning via MLflow.

MLflow Tracing samlar in detaljerad körningsinformation för gen AI-program. Den loggar indata, utdata och metadata för varje steg, vilket hjälper dig att felsöka problem och analysera prestanda.

När du använder UCFunctionToolkitgenererar hämtningsverktygen automatiskt RETRIEVER intervalltyper i MLflow-loggarna om deras utdata matchar MLflow retriever-schemat. Se MLflow Tracing Schema.

För mer information om UCFunctionToolkit, se dokumentationen för Unity Catalog .

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

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

• SQL-klienter kan begränsa det maximala antalet rader eller byte som returneras. För att förhindra datatrunkering bör du trunkera kolumnvärden som returneras av UDF. Du kan till exempel använda substring(chunked_text, 0, 8192) för att minska storleken på stora innehållskolumner och undvika att rader trunkeras under utförandet.
• 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

För att följa MLflow retriever-schemat bör funktionen retriever returnera en dokumenttyp och använda fältet metadata i dokumentklassen 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)

Ange 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 kolumnnamnen i den returnerade tabellen 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 kolumner i hämtarens schema genom att ange en lista med kolumnnamn med fältet other_columns.

Om du har flera retrievers kan du definiera flera scheman med hjälp av unika namn för varje retrieverschema.

Datahämtningsschemauppsättningen när agenten skapas påverkar underordnade program och arbetsflöden, till exempel granskningsappen och utvärderingsuppsättningarna. Mer specifikt fungerar kolumnen doc_uri som primär identifierare 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 begränsa 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 funktionsagentverktyg för Unity Catalog lägger du till verktyget i en AI-agent. Se Lägg till verktyg för Unity Catalog i agenter.