Condividi tramite

Connettere agenti a dati non strutturati

Gli agenti di intelligenza artificiale spesso devono eseguire query su dati non strutturati come raccolte di documenti, knowledge base o corpora di testo per rispondere alle domande e fornire risposte in grado di essere consapevoli del contesto.

Databricks offre più approcci per la connessione degli agenti a dati non strutturati negli indici di Ricerca vettoriale e negli archivi vettoriali esterni. Usare i server MCP preconfigurato per l'accesso immediato agli indici di Ricerca vettoriale di Databricks, sviluppare strumenti di recupero in locale con pacchetti AI Bridge o creare funzioni di recupero personalizzate per flussi di lavoro specializzati.

Questa pagina illustra come:

Eseguire query su un indice di ricerca vettoriale di Databricks

Se l'agente deve eseguire query su un indice di ricerca vettoriale di Databricks, usare il server MCP gestito da Databricks per questo tipo:

  1. Creare un indice di ricerca vettoriale usando incorporamenti gestiti di Databricks. Vedere Creare endpoint e indici di ricerca vettoriali.
  2. Creare un agente e connetterlo all'URL MCP gestito preconfigurato per l'indice di ricerca vettoriale: https://<workspace-hostname>/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}.

Per informazioni su come creare un agente che si connette ai server MCP gestiti, vedere Usare i server MCP gestiti di Databricks.

Sviluppare localmente gli strumenti di recupero di ricerca vettoriale con AI Bridge

Per creare localmente uno strumento di recupero per la ricerca vettoriale di Databricks, utilizzare i pacchetti Databricks AI Bridge come databricks-langchain e databricks-openai. Questi pacchetti includono funzioni helper come from_vector_search e from_uc_function per creare recuperatori da risorse di Databricks esistenti.

LangChain/LangGraph

Installare la versione più recente di databricks-langchain che include Databricks AI Bridge.

%pip install --upgrade databricks-langchain

Il codice seguente crea prototipi di uno strumento di recupero che esegue una query su un indice ipotetico di ricerca vettoriale e lo associa a un LLM in locale, in modo da poter testare il comportamento di chiamata allo strumento.

Fornire un descrittivo tool_description per aiutare l'agente a comprendere lo strumento e determinare quando richiamarlo.

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

Per gli scenari che usano indici di accesso diretto o indici di Sincronizzazione Delta tramite incorporamenti autogestiti, è necessario configurare VectorSearchRetrieverTool e specificare un modello di incorporamento personalizzato e una colonna di testo. Consultare le opzioni per fornire incorporazioni.

L'esempio seguente mostra come configurare un VectorSearchRetrieverTool con le chiavi columns e embedding.

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

Per altri dettagli, vedere la documentazione dell'API per VectorSearchRetrieverTool.

OpenAI

Installare la versione più recente di databricks-openai che include Databricks AI Bridge.

%pip install --upgrade databricks-openai

Il codice seguente crea prototipi di un retriever che esegue una query su un indice ipotetico di ricerca vettoriale e lo integra con i modelli GPT di OpenAI.

Fornire un descrittivo tool_description per aiutare l'agente a comprendere lo strumento e determinare quando richiamarlo.

Per altre informazioni sulle raccomandazioni openAI per gli strumenti, vedere documentazione relativa alla chiamata di funzioni OpenAI.

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

Per gli scenari che usano indici di accesso diretto o indici di Sincronizzazione Delta tramite incorporamenti autogestiti, è necessario configurare VectorSearchRetrieverTool e specificare un modello di incorporamento personalizzato e una colonna di testo. Consultare le opzioni per fornire incorporazioni.

L'esempio seguente mostra come configurare un VectorSearchRetrieverTool con le chiavi columns e embedding.

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

Per altri dettagli, vedere la documentazione dell'API per VectorSearchRetrieverTool.

Quando lo strumento locale è pronto, è possibile crearlo direttamente in produzione come parte del codice dell'agente o eseguirne la migrazione a una funzione del catalogo Unity, che offre una migliore individuabilità e governance, ma presenta alcune limitazioni.

La sezione seguente illustra come eseguire la migrazione del retriever a una funzione del catalogo Unity.

Eseguire una query su un indice vettoriale ospitato all'esterno di Databricks

Se l'indice vettoriale è ospitato all'esterno di Azure Databricks, è possibile creare una connessione di Unity Catalog per connettersi al servizio esterno e usare la connessione nel codice dell'agente. Vedere Connettere gli strumenti dell'agente di intelligenza artificiale ai servizi esterni.

Nell'esempio seguente viene creato un retriever che chiama un indice vettoriale ospitato all'esterno di Databricks per un agente PyFunc.

  1. Creare una connessione del catalogo Unity al servizio esterno, in questo caso Azure.

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

  2. Definire lo strumento retriever nel codice dell'agente usando la connessione al catalogo Unity. Questo esempio usa i decoratori MLflow per abilitare il tracciamento dell'agente.

    Nota

    Per essere conforme allo schema del retriever MLflow, la funzione retriever deve restituire un List[Document] oggetto e utilizzare il metadata campo nella classe Document per aggiungere altri attributi al documento restituito, ad esempio doc_uri e similarity_score. Vedere Documento MLflow.

    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. Per eseguire il retriever, eseguire il codice Python seguente. Facoltativamente, è possibile includere filtri di ricerca vettoriale nella richiesta per filtrare i risultati.

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

Aggiungere il tracciamento a un retriever

Aggiungi il tracciamento MLflow per monitorare e eseguire il debug del retriever. La traccia consente di visualizzare input, output e metadati per ogni passaggio di esecuzione.

Nell'esempio precedente, il decorator @mlflow.trace viene aggiunto sia ai __call__ metodi che a quelli di analisi. Il decoratore crea un intervallo che inizia quando la funzione viene richiamata e termina al suo ritorno. MLflow registra automaticamente l'input e l'output della funzione ed eventuali eccezioni generate.

Nota

Gli utenti della libreria LangChain, LlamaIndex e OpenAI possono usare la registrazione automatica di MLflow oltre a definire manualmente le tracce con l'elemento decorator. Vedere Aggiungere tracce alle applicazioni: traccia automatica e manuale.

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

Per garantire che le applicazioni downstream come Agent Evaluation e AI Playground eseguano correttamente la rappresentazione della traccia del retriever, assicurarsi che il decoratore soddisfi i seguenti requisiti:

Impostare lo schema del retriever per garantire la compatibilità di MLflow

Se la traccia restituita dal retriever o span_type="RETRIEVER" non è conforme allo schema di recupero standard di MLflow, è necessario eseguire manualmente il mapping dello schema restituito ai campi previsti di MLflow. In questo modo, MLflow può tracciare correttamente il retriever e visualizzare le tracce nelle applicazioni successive.

Per impostare manualmente lo schema del retriever:

  1. Chiamare mlflow.models.set_retriever_schema quando si definisce l'agente. Usare set_retriever_schema per mappare i nomi delle colonne nella tabella restituita ai campi previsti di MLflow, ad esempio primary_key, text_column e 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. Specificare colonne aggiuntive nello schema del retriever fornendo un elenco di nomi di colonna con il other_columns campo .

  3. Se si dispone di più retriever, è possibile definire più schemi usando nomi univoci per ogni schema di recupero.

Lo schema del retriever impostato durante la creazione dell'agente influisce su applicazioni e flussi di lavoro downstream, ad esempio l'app di revisione e i set di valutazione. In particolare, la colonna doc_uri funge da identificatore primario per i documenti restituiti dal retriever.

  • L'app di revisione visualizza il doc_uri per aiutare i revisori a valutare le risposte e tracciare le origini dei documenti. Vedi Revisione dell'interfaccia utente dell'app.
  • I set di valutazione utilizzano doc_uri per confrontare i risultati del sistema di recupero con i set di dati di valutazione predefiniti, al fine di determinare il recupero e la precisione del sistema. Vedere Set di valutazione (MLflow 2).

strumento di ricerca vettoriale con funzioni del Catalogo Unity

Nota

Databricks consiglia i server MCP per la maggior parte degli strumenti dell'agente, ma la definizione degli strumenti con le funzioni di Unity Catalog rimane disponibile per la creazione di prototipi.

È possibile creare una funzione di Unity Catalog che incapsula una query sull'indice di ricerca a vettori Mosaic AI. Questo approccio:

  • Supporta i casi d'uso di produzione con governance e facilità di individuazione
  • Usa la funzione SQL vector_search() sotto le quinte
  • Supporta il tracciamento automatico di MLflow
    • È necessario allineare l'output della funzione allo schema del retriever MLflow usando gli page_content alias e metadata .
    • Tutte le colonne di metadati aggiuntive devono essere aggiunte alla metadata colonna usando la funzione di mapping SQL, anziché come chiavi di output di primo livello.

Eseguire il codice seguente in un notebook o in un editor SQL per creare la funzione:

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
  )

Per usare questo strumento di recupero nel tuo agente di intelligenza artificiale, racchiudilo con UCFunctionToolkit. Questo consente la tracciatura automatica tramite MLflow, generando automaticamente tipi di span nei log di MLflow.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Gli strumenti di recupero del catalogo Unity presentano le avvertenze seguenti:

  • I client SQL possono limitare il numero massimo di righe o byte restituiti. Per evitare il troncamento dei dati, è necessario troncare i valori delle colonne restituiti dalla UDF (funzione definita dall'utente). Ad esempio, è possibile usare substring(chunked_text, 0, 8192) per ridurre le dimensioni delle colonne di contenuto di grandi dimensioni ed evitare il troncamento delle righe durante l'esecuzione.
  • Poiché questo strumento è un wrapper per la funzione vector_search(), è soggetto alle stesse limitazioni della funzione vector_search(). Vedere Limitazioni.

Per altre informazioni sui UCFunctionToolkit, vedere documentazione di Unity Catalog.

  • Last updated on