Hulpprogramma's voor het bouwen en traceren van ongestructureerde dataretrievers

Gebruik het Mozaïek AI Agent Framework om hulpprogramma's te bouwen waarmee AI-agents ongestructureerde gegevens kunnen opvragen, zoals een verzameling documenten. Op deze pagina ziet u hoe u het volgende kunt doen:

• Lokaal retrievers ontwikkelen
• Een retriever maken met Unity Catalog-functies
• Externe vectorindexen opvragen
• MLflow-tracering toevoegen voor waarneembaarheid

Zie AI-agenthulpprogramma'svoor meer informatie over agenthulpprogramma's.

Hulpprogramma's voor vectorzoekopdrachten lokaal ontwikkelen met AI Bridge

De snelste manier om te beginnen met het bouwen van een Databricks Vector Search-retriever-hulpprogramma is het lokaal ontwikkelen en testen met behulp van Databricks AI Bridge-pakketten zoals databricks-langchain en databricks-openai.

LangChain/LangGraph

Installeer de nieuwste versie van databricks-langchain die Databricks AI Bridge bevat.

%pip install --upgrade databricks-langchain

De volgende code maakt een prototype van een retriever-tool die een query uitvoert op een hypothetische vectorzoekindex en deze lokaal aan een LLM koppelt, zodat u het gedrag van het aanroepen van tools kunt testen.

Geef een beschrijvende tool_description functie op om de agent te helpen het hulpprogramma te begrijpen en te bepalen wanneer het moet worden aangeroepen.

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-sonnet-4-5")
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?")

Bij scenario's die gebruikmaken van indexen voor directe toegang of Delta Sync-indexen met zelfbeheer inbeddingen, moet u het VectorSearchRetrieverTool configureren en een aangepast inbeddingsmodel en tekstkolom specificeren. Zie opties voor het bieden van embeddingen.

In het volgende voorbeeld ziet u hoe u een VectorSearchRetrieverTool met columns en embedding sleutels configureert.

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

Zie de API-documenten voor meer informatie.VectorSearchRetrieverTool

OpenAI

Installeer de nieuwste versie van databricks-openai die Databricks AI Bridge bevat.

%pip install --upgrade databricks-openai

De volgende code prototypeert een retriever die een query uitvoert op een hypothetische vectorzoekindex en deze integreert met GPT-modellen van OpenAI.

Geef een beschrijvende tool_description functie op om de agent te helpen het hulpprogramma te begrijpen en te bepalen wanneer het moet worden aangeroepen.

Zie de documentatie OpenAI-functieaanroepenvoor meer informatie over OpenAI-aanbevelingen voor hulpprogramma's.

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

Bij scenario's die gebruikmaken van indexen voor directe toegang of Delta Sync-indexen met zelfbeheer inbeddingen, moet u het VectorSearchRetrieverTool configureren en een aangepast inbeddingsmodel en tekstkolom specificeren. Zie opties voor het bieden van embeddingen.

In het volgende voorbeeld ziet u hoe u een VectorSearchRetrieverTool met columns en embedding sleutels configureert.

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

Zie de API-documenten voor meer informatie.VectorSearchRetrieverTool

Zodra uw lokale hulpprogramma gereed is, kunt u het rechtstreeks produceren als onderdeel van uw agentcode of migreren naar een Unity Catalog-functie, die betere detectie en governance biedt, maar bepaalde beperkingen heeft.

In de volgende sectie ziet u hoe u de retriever migreert naar een Unity Catalog-functie.

Vector Search Retriever-hulpprogramma met Unity Catalog-functies

U kunt een Unity Catalog-functie maken waarmee een Mozaïek AI Vector Search-indexquery wordt verpakt. Deze aanpak:

• Ondersteunt toepassingsgevallen voor productie met governance en ontdekbaarheid
• Maakt gebruik van de SQL-functie vector_search() achter de schermen
• Ondersteunt automatische MLflow-tracering
  • U moet de uitvoer van de functie uitlijnen op het MLflow retriever-schema met behulp van de page_content en metadata aliassen.
  • Eventuele extra metagegevenskolommen moeten worden toegevoegd aan de metadata kolom met behulp van de SQL-toewijzingsfunctie in plaats van als uitvoersleutels op het hoogste niveau.

Voer de volgende code uit in een notebook of SQL-editor om de functie te maken:

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
  )

Als u dit retriever-hulpprogramma in uw AI-agent wilt gebruiken, verpakt u deze met UCFunctionToolkit. Dit maakt automatische tracering via MLflow mogelijk door automatisch spantypen in MLflow-logboeken te genereren RETRIEVER .

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Hulpprogramma's voor het ophalen van Unity Catalog hebben de volgende opmerkingen:

• SQL-clients kunnen het maximum aantal geretourneerde rijen of bytes beperken. Om te voorkomen dat gegevens worden afgekort, moet u de kolomwaarden inkorten die door de UDF worden geretourneerd. U kunt bijvoorbeeld substring(chunked_text, 0, 8192) gebruiken om de grootte van kolommen met grote inhoud te verkleinen en afkapping van rijen te voorkomen tijdens de uitvoering.
• Omdat dit hulpprogramma een wrapper is voor de vector_search() functie, gelden dezelfde beperkingen als de functie vector_search(). Zie Beperkingen.

Zie de documentatie van UCFunctionToolkitvoor meer informatie over .

Retriever die een query uitvoert op een vectorindex die buiten Databricks wordt gehost

Als uw vectorindex buiten Azure Databricks wordt gehost, kunt u een Unity Catalog-verbinding maken om verbinding te maken met de externe service en de verbinding in uw agentcode te gebruiken. Zie Ai-agenthulpprogramma's verbinden met externe services.

In het volgende voorbeeld wordt een retriever gemaakt die een vector index aanroept, gehost buiten Databricks, voor een PyFunc-gebaseerde agent.

1. Maak een Unity Catalog-verbinding met de externe service, in dit geval Azure.

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

2. Definieer het hulpprogramma retriever in agentcode met behulp van de Unity Catalog-verbinding. In dit voorbeeld worden MLflow-decorators gebruikt om agenttracering in te schakelen.

   Notitie

   Om te voldoen aan het MLflow retriever-schema, moet de functie retriever een List[Document] object retourneren en het metadata veld in de documentklasse gebruiken om extra kenmerken toe te voegen aan het geretourneerde document, zoals doc_uri en similarity_score. Zie MLflow-document.

   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. Voer de volgende Python-code uit om de retriever uit te voeren. U kunt desgewenst Vector Search-filters opnemen in de aanvraag om resultaten te filteren.

   retriever = VectorSearchRetriever()
   query = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
   results = retriever(query, score_threshold=0.1)
   

Tracering toevoegen aan een retriever

Voeg MLflow-tracering toe om uw retriever te bewaken en fouten op te sporen. Met tracering kunt u invoer, uitvoer en metagegevens bekijken voor elke stap van de uitvoering.

In het vorige voorbeeld wordt de @mlflow.trace decorator toegevoegd aan zowel de __call__ als de parseringsmethoden. De decorator maakt een scope die begint wanneer de functie wordt aangeroepen en eindigt wanneer deze terugkeert. MLflow registreert automatisch de invoer en uitvoer van de functie en eventuele uitzonderingen die zijn gegenereerd.

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]:
    ...

Om ervoor te zorgen dat toepassingen later in het proces, zoals evaluatie van agenten en de AI Playground, de retriever-tracering correct weergeven, moet u erop letten dat de decorator aan de volgende vereisten voldoet:

• Gebruik (https://mlflow.org/docs/latest/tracing/tracing-schema.html#retriever-spans) en zorg ervoor dat de functie een List[Document]-object retourneert.
• De traceringsnaam en de retriever_schema naam moeten overeenkomen om de tracering correct te configureren. Zie de volgende sectie voor meer informatie over het instellen van het retriever-schema.

Ophalenschema instellen om compatibiliteit met MLflow te garanderen

Als de tracering die wordt geretourneerd door de retriever of span_type="RETRIEVER" niet voldoet aan het standaard-retrieverschema van MLflow, moet u het geretourneerde schema handmatig toewijzen aan de verwachte velden van MLflow. Dit zorgt ervoor dat MLflow uw retriever correct kan traceren en traceringen kan weergeven in downstreamtoepassingen.

Het retriever-schema handmatig instellen:

1. Roep mlflow.models.set_retriever_schema aan wanneer u uw agent definieert. Gebruik set_retriever_schema om de kolomnamen in de geretourneerde tabel toe te wijzen aan de verwachte velden van MLflow, zoals primary_key, text_column, en 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. Geef extra kolommen op in het schema van de retriever door een lijst met kolomnamen met het other_columns veld op te geven.

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

Het retriever-schema dat is ingesteld tijdens het maken van agents, beïnvloedt downstream toepassingen en processen, zoals de beoordelingsapp en evaluatiesets. De kolom doc_uri fungeert met name als de primaire id voor documenten die door de retriever worden geretourneerd.

• In de beoordelings-app wordt de doc_uri weergegeven om revisoren te helpen antwoorden te beoordelen en documentoorsprongen te traceren. Zie App-UI beoordelen.
• Evaluatiesets gebruiken doc_uri om de retrieverresultaten te vergelijken met vooraf gedefinieerde evaluatiesets en de recall en precisie van de retriever te bepalen. Zie Evaluatiesets (MLflow 2).

Volgende stappen

Nadat u een Unity Catalog-hulpprogramma hebt gemaakt, voegt u dit toe aan uw agent. Zie Een agenthulpprogramma maken.