AI-agenten en -hulpmiddelen voor het ophalen van ongestructureerde informatie

Artikel
03/09/2025

In dit artikel wordt beschreven hoe u ai-agenthulpprogramma's maakt voor het ophalen van ongestructureerde gegevens met behulp van het Mozaïek AI Agent Framework. Ongestructureerde retrievers stellen agents in staat om een query uit te voeren op niet-gestructureerde gegevensbronnen, zoals een documentinhoud, met behulp van vectorzoekindexen.

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

Hulpprogramma's voor vectorzoekopdrachten lokaal ontwikkelen met AI Bridge

De eenvoudigste manier om een Databricks Vector Search-retriever-hulpprogramma te ontwikkelen, is lokaal. Gebruik Databricks AI Bridge-pakketten zoals databricks-langchain en databricks-openai om snel ophaalmogelijkheden toe te voegen aan een agent en te experimenteren met queryparameters. Deze benadering maakt een snelle iteratie mogelijk tijdens de eerste ontwikkeling.

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. Zie Vector Search Retriever-hulpprogramma met Unity Catalog-functies.

Notitie

Als u een externe vectorindex wilt gebruiken die buiten Databricks wordt gehost, raadpleegt u Vector Search-retriever met behulp van een vectorindex die buiten Databrickswordt gehost.

LangChain/LangGraph

Met de volgende code wordt een retriever-hulpprogramma geprotoaliseerd en wordt het lokaal aan een LLM gekoppeld, zodat u met de agent kunt chatten om het gedrag van het aanroepen van hulpprogramma's te testen.

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

%pip install --upgrade databricks-langchain

In het volgende voorbeeld wordt een query uitgevoerd op een hypothetische vectorzoekindex waarmee inhoud wordt opgehaald uit de productdocumentatie van Databricks.

Geef een duidelijke en beschrijvende tool_description. De agent LLM gebruikt de tool_description om inzicht te krijgen in het hulpprogramma en te bepalen wanneer het hulpprogramma 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-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?")

Notitie

Bij het initialiseren van de VectorSearchRetrieverToolzijn de argumenten text_column en embedding vereist voor Delta Sync-indexen met zelfbeheerde insluitingen en Direct Vector Access-indexen. Zie opties voor het bieden van embeddingen.

Zie de API-documenten voor VectorSearchRetrieverTool voor meer informatie

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

OpenAI

De volgende code prototypeert een vector zoekophaaldienst en integreert deze met de GPT-modellen van OpenAI.

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

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

%pip install --upgrade databricks-openai

In het volgende voorbeeld wordt een query uitgevoerd op een hypothetische vectorzoekindex waarmee inhoud wordt opgehaald uit de productdocumentatie van Databricks.

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

Notitie

Zie de API-documenten voor VectorSearchRetrieverTool voor meer informatie

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

Vector Search Retriever-hulpprogramma met Unity Catalog-functies

In het volgende voorbeeld wordt een ophaalprogramma gemaakt met behulp van een Unity Catalog-functie om gegevens op te vragen uit een Mozaïek AI Vector Search-index.

De functie Unity Catalog databricks_docs_vector_search voert een query uit op een hypothetische Vector Search-index met Databricks-documentatie. Met deze functie wordt de Databricks SQL-functie vector_search() verpakt en wordt de uitvoer uitgelijnd met het MLflow retriever-schema. met behulp van de page_content en metadata aliassen.

Notitie

Om te voldoen aan het MLflow retriever-schema, moeten eventuele extra metagegevenskolommen 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.

MLflow Tracing legt gedetailleerde uitvoeringsinformatie vast voor gen AI-toepassingen. Het registreert invoer, uitvoer en metagegevens voor elke stap, zodat u problemen kunt opsporen en prestaties kunt analyseren.

Wanneer u UCFunctionToolkitgebruikt, genereren retrievers automatisch RETRIEVER spantypen in MLflow-logboeken als hun uitvoer voldoet aan het schema van de MLflow-retriever. Zie schema voor MLflow-tracering.

Zie de documentatie van UCFunctionToolkitvoor meer informatie over .

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Dit ophaalhulpmiddel heeft de volgende kanttekeningen:

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.

Vector Search-retriever met behulp van 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 verbindingsagentcode te gebruiken. Zie Ai-agenthulpprogramma's verbinden met externe servicesvoor meer informatie.

In het volgende voorbeeld wordt een Vector Search-retriever gemaakt die een vectorindex aanroept die buiten Databricks wordt gehost voor een PyFunc-agent.

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>')
);

Definieer het hulpprogramma retriever in agentcode met behulp van de Unity Catalog Connection die u hebt gemaakt. In dit voorbeeld worden MLflow-decorators gebruikt om agenttracering in te schakelen.

Notitie

Als u wilt voldoen aan het schema van de MLflow-retriever, moet de functie retriever een documenttype retourneren en het metadata veld in de documentklasse gebruiken om extra kenmerken toe te voegen aan het geretourneerde document, zoals like doc_uri en similarity_score.

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

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)

Schema voor ophalen instellen

Als de trace die wordt geretourneerd door de retriever of span_type="RETRIEVER" niet overeenkomt met 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 de traceringen correct kan weergeven in downstreamtoepassingen.

Als u het retriever-schema handmatig wilt instellen, roept u 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_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"],
)

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

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 worden gebruikt om de resultaten van de retriever te vergelijken met vooraf gedefinieerde evaluatiesets om de herinnering en precisie van de retriever te bepalen. Zie Evaluatiesets.

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 snel de bron van bugs en onverwacht gedrag kunt vaststellen.

In dit voorbeeld wordt de @mlflow.trace decorator gebruikt om een spoor te creëren voor de retriever en parser. Zie Agent-waarneembaarheid met MLflow Tracingvoor andere opties voor het instellen van traceringsmethoden.

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.

Notitie

Gebruikers van de LangChain-, LlamaIndex- en OpenAI-bibliotheek kunnen automatische logboekregistratie van MLflow gebruiken in plaats van traceringen handmatig te definiëren met de decorator. Zie Automatische tracering.

...
@mlflow.trace(span_type="RETRIEVER", name="vector_search")
def __call__(self, query: str) -> 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 span_type="RETRIEVER" en zorg ervoor dat de functie List[Document] object retourneert. Zie Retriever reikwijdtes.
De traceringsnaam en de retriever_schema naam moeten overeenkomen om de tracering correct te configureren.

Volgende stappen

Nadat u een hulpprogramma voor de Unity Catalog-functieagent hebt gemaakt, voegt u het hulpprogramma toe aan een AI-agent. Zie Voeg Unity Catalog-hulpprogramma's toe aan agents.

Delen via

AI-agenten en -hulpmiddelen voor het ophalen van ongestructureerde informatie

Hulpprogramma's voor vectorzoekopdrachten lokaal ontwikkelen met AI Bridge

LangChain/LangGraph

OpenAI

Vector Search Retriever-hulpprogramma met Unity Catalog-functies

Vector Search-retriever met behulp van een vectorindex die buiten Databricks wordt gehost

Schema voor ophalen instellen

De retriever traceren

Volgende stappen

Feedback

Aanvullende resources