Skapa och spåra hämtningsverktyg för ostrukturerade data

Använd Mosaic AI Agent Framework för att skapa verktyg som gör att AI-agenter kan köra frågor mot ostrukturerade data, till exempel en samling dokument. Den här sidan visar hur du:

• Utveckla hämtningar lokalt
• Skapa en retriever med hjälp av Unity Catalog-funktioner
• Sök i externa vektorindex
• Lägga till MLflow-spårning för observerbarhet

Mer information om agentverktyg finns i AI-agentverktyg.

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

Det snabbaste sättet att börja skapa ett Databricks Vector Search-hämtningsverktyg är att utveckla och testa det lokalt med databricks AI Bridge-paket som databricks-langchain och databricks-openai.

LangChain/LangGraph

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

%pip install --upgrade databricks-langchain

Följande kod prototypar ett retriever-verktyg som utfrågar ett hypotetiskt vektorsökindex och binder det till LLM:en lokalt så att du kan testa dess beteende vid verktygsanrop.

Ange ett beskrivande tool_description för att hjälpa agenten att förstå verktyget och avgöra när det ska anropas.

from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool.
vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation."
)

# 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-claude-3-7-sonnet")
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?")

För scenarier som använder antingen direktåtkomstindex eller Delta Sync-index med självhanterade inbäddningar måste du konfigurera VectorSearchRetrieverTool och ange en anpassad inbäddningsmodell och textkolumn. Se alternativ för att tillhandahålla inbäddningar.

I följande exempel visas hur du konfigurerar en VectorSearchRetrieverTool med columns och embedding nycklar.

from databricks_langchain import VectorSearchRetrieverTool
from databricks_langchain import DatabricksEmbeddings

embedding_model = DatabricksEmbeddings(
    endpoint="databricks-bge-large-en",
)

vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
  num_results=5, # Max number of documents to return
  columns=["primary_key", "text_column"], # List of columns to include in the search
  filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
  query_type="ANN", # Query type ("ANN" or "HYBRID").
  tool_name="name of the tool", # Used by the LLM to understand the purpose of the tool
  tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
  text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
  embedding=embedding_model # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Mer information finns i API-dokumenten för VectorSearchRetrieverTool.

OpenAI (på engelska)

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

%pip install --upgrade databricks-openai

Följande kod skapar en prototyp av en retriever som gör förfrågningar till ett vektorbaserat sökindex och integreras med OpenAI:s GPT-modeller.

Ange ett beskrivande tool_description för att hjälpa agenten att förstå verktyget och avgöra när det ska anropas.

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

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

# Initialize OpenAI client
client = OpenAI(api_key=<your_API_key>)

# Initialize the retriever tool
dbvs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation"
)

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

För scenarier som använder antingen direktåtkomstindex eller Delta Sync-index med självhanterade inbäddningar måste du konfigurera VectorSearchRetrieverTool och ange en anpassad inbäddningsmodell och textkolumn. Se alternativ för att tillhandahålla inbäddningar.

I följande exempel visas hur du konfigurerar en VectorSearchRetrieverTool med columns och embedding nycklar.

from databricks_openai import VectorSearchRetrieverTool

vs_tool = VectorSearchRetrieverTool(
    index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
    num_results=5, # Max number of documents to return
    columns=["primary_key", "text_column"], # List of columns to include in the search
    filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
    query_type="ANN", # Query type ("ANN" or "HYBRID").
    tool_name="name of the tool", # Used by the LLM to understand the purpose of the tool
    tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
    text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
    embedding_model_name="databricks-bge-large-en" # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Mer information finns i API-dokumenten för VectorSearchRetrieverTool.

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.

I följande avsnitt visas hur du migrerar hämtaren till en Unity Catalog-funktion.

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

Du kan skapa en Unity Catalog-funktion som omsluter en Mosaic AI Vector Search-indexfråga . Den här metoden:

• Stöder produktionsanvändningsfall med styrning och identifiering
• Använder sql-funktionen vector_search() i bakgrunden
• Stöder automatisk MLflow-spårning
  • Du måste justera funktionens utdata till MLflow retriever-schemat med hjälp av aliasen page_content och metadata .
  • Eventuella ytterligare metadatakolumner måste läggas till i metadata kolumnen med hjälp av SQL Map-funktionen i 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 genom att automatiskt generera RETRIEVER intervalltyper i MLflow-loggar.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Unity Catalog retriever-verktyg 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.

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

Retriever som frågar ett vektorindex som finns utanför Databricks

Om ditt vektorindex finns utanför Azure Databricks kan du skapa en Unity Catalog-anslutning för att ansluta till den externa tjänsten och använda anslutningen i agentkoden. Se Anslut AI-agentverktyg till externa tjänster.

Det följande exemplet visar hur man skapar en retriever som anropar ett vektorindex som finns utanför Databricks för en PyFunc-agent.

1. Skapa en Unity-kataloganslutning till den externa tjänsten, i det här fallet Azure.

   CREATE CONNECTION ${connection_name}
   TYPE HTTP
   OPTIONS (
     host 'https://example.search.windows.net',
     base_path '/',
     bearer_token secret ('<secret-scope>','<secret-key>')
   );
   

2. Definiera verktyget retriever i agentkoden med hjälp av Unity Catalog-anslutningen. I det här exemplet används MLflow-dekoratörer för att aktivera agentspårning.

   Anteckning

   För att följa MLflow retriever-schemat ska funktionen retriever returnera ett List[Document] objekt och använda metadata fältet i klassen Dokument för att lägga till ytterligare attribut i det returnerade dokumentet, till exempel doc_uri och similarity_score. Se MLflow-dokument.

   import mlflow
   import json
   
   from mlflow.entities import Document
   from typing import List, Dict, Any
   from dataclasses import asdict
   
   class VectorSearchRetriever:
     """
     Class using Databricks Vector Search to retrieve relevant documents.
     """
   
     def __init__(self):
       self.azure_search_index = "hotels_vector_index"
   
     @mlflow.trace(span_type="RETRIEVER", name="vector_search")
     def __call__(self, query_vector: List[Any], score_threshold=None) -> List[Document]:
       """
       Performs vector search to retrieve relevant chunks.
       Args:
         query: Search query.
         score_threshold: Score threshold to use for the query.
   
       Returns:
         List of retrieved Documents.
       """
       from databricks.sdk import WorkspaceClient
       from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod
   
       json = {
         "count": true,
         "select": "HotelId, HotelName, Description, Category",
         "vectorQueries": [
           {
             "vector": query_vector,
             "k": 7,
             "fields": "DescriptionVector",
             "kind": "vector",
             "exhaustive": true,
           }
         ],
       }
   
       response = (
         WorkspaceClient()
         .serving_endpoints.http_request(
           conn=connection_name,
           method=ExternalFunctionRequestHttpMethod.POST,
           path=f"indexes/{self.azure_search_index}/docs/search?api-version=2023-07-01-Preview",
           json=json,
         )
         .text
       )
   
       documents = self.convert_vector_search_to_documents(response, 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 = []
   
       for item in vs_results.get("value", []):
         score = item.get("@search.score", 0)
   
         if score >= score_threshold:
           metadata = {
             "score": score,
             "HotelName": item.get("HotelName"),
             "Category": item.get("Category"),
           }
   
           doc = Document(
             page_content=item.get("Description", ""),
             metadata=metadata,
             id=item.get("HotelId"),
           )
           docs.append(doc)
   
       return docs
   

3. 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 = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
   results = retriever(query, score_threshold=0.1)
   

Lägga till spårning i en retriever

Lägg till MLflow-spårning för att övervaka och felsöka din retriever. Med spårning kan du visa indata, utdata och metadata för varje körningssteg.

I föregående exempel läggs @mlflow.trace-dekoratören till både __call__ metoderna och parsning. 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.

import mlflow
from mlflow.entities import Document

## This code snippet has been truncated for brevity, see the full retriever example above
class VectorSearchRetriever:
  ...

  # Create a RETRIEVER span. The span name must match the retriever schema name.
  @mlflow.trace(span_type="RETRIEVER", name="vector_search")
  def __call__(...) -> List[Document]:
    ...

  # Create a PARSER span.
  @mlflow.trace(span_type="PARSER")
  def parse_results(...) -> 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 (https://mlflow.org/docs/latest/tracing/tracing-schema.html#retriever-spans) och se till att funktionen returnerar ett List[Document]-objekt.
• Spårningsnamnet och retriever_schema namnet måste matcha för att kunna konfigurera spårningen korrekt. Se följande avsnitt för att lära dig hur du ställer in retriever-schemat.

Ange retriever-schema för att säkerställa MLflow-kompatibilitet

Om spårningen som returneras från hämtaren eller span_type="RETRIEVER" inte överensstämmer med MLflows standardåterhämtningsschema måste du mappa det returnerade schemat manuellt till MLflows förväntade fält. Detta säkerställer att MLflow kan korrekt spåra din retriever och visa spår i senare applikationer.

Så här anger du upphämtningsschema manuellt:

1. Anropa 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"],
   )
   

2. Ange ytterligare kolumner i hämtarens schema genom att ange en lista med kolumnnamn med fältet other_columns .

3. 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 över Granska appens användargränssnitt.
• Utvärderingsuppsättningar används doc_uri för att jämföra hämtningsresultat med fördefinierade utvärderingsdatauppsättningar för att fastställa hämtarens träffsäkerhet och precision. Se Utvärderingsuppsättningar (MLflow 2).

Nästa steg

När du har skapat ett Unity Catalog-verktyg lägger du till det i din agent. Se Skapa ett agentverktyg.