Connettere agenti a dati non strutturati

Eseguire query su un indice di ricerca vettoriale di Databricks usando il mcp di ricerca vettoriale

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.

Eseguire query su un indice di ricerca 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 recuperatore che chiama un indice vettoriale ospitato all'esterno di Databricks per un agente basato su 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)
   

Sviluppare un retriever localmente usando 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.

Dopo che 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.

Eseguire query su Vector Search di Databricks usando le funzioni UC (deprecate)

È 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 potrebbero limitare il numero massimo di righe o byte restituiti. Per evitare il troncamento dei dati, troncate i valori delle colonne restituiti dalla UDF. 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 vector_search() funzione, è soggetto alle stesse limitazioni della vector_search() funzione. Vedere Limitazioni.

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