Pular para o conteúdo principal

Conectar agentes a dados não estruturados

Agentes de AI frequentemente precisam query dados não estruturados, como coleções de documentos, bases de conhecimento ou corpus de texto, para responder a perguntas e fornecer respostas contextualizadas.

Databricks oferece várias abordagens para conectar agentes a dados não estruturados em índices de Pesquisa AI e armazenamentos de vetor externos. Use servidores MCP pré-configurados para acesso imediato a índices de Pesquisa Databricks AI, desenvolva ferramentas de recuperação localmente com pacotes do AI Bridge ou crie funções de recuperação personalizadas para fluxos de trabalho especializados.

Databricks AI Search era anteriormente conhecido como Databricks Vector Search. O prefixo de URL legado /api/2.0/mcp/vector-search/ continua a funcionar para compatibilidade com versões anteriores.

Fazer query em um índice de Pesquisa de Databricks AI usando MCP

Use o servidor MCP de Pesquisa de Databricks AI gerenciado pela Databricks para dar ao seu agente acesso a um índice de Databricks AI. Primeiro, crie um índice usando incorporações gerenciadas pela Databricks. Consulte Criar Endpoint e índices de Pesquisa de IA.

A URL MCP gerenciada para pesquisa de AI é https://<workspace-hostname>/api/2.0/mcp/ai-search/{catalog}/{schema}/{index_name}. Conecte-se a ele e liste as ferramentas que ele expõe:

Python
from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksMCPClient

workspace_client = WorkspaceClient()
host = workspace_client.config.host

mcp_client = DatabricksMCPClient(
server_url=f"{host}/api/2.0/mcp/ai-search/<catalog>/<schema>/<index-name>",
workspace_client=workspace_client,
)
tools = mcp_client.list_tools()

Para construir e implantar um agente que usa este servidor, consulte Usar servidores MCP em agentes. Conceda ao agente SELECT no objeto protegível do Unity Catalog do índice.

Outras abordagens

Consultar um índice de pesquisa vetorial fora do Databricks

query um índice de pesquisa vetorial hospedado fora do Databricks

Se o seu índice vetorial estiver hospedado fora do Databricks, você poderá criar uma conexão do Unity Catalog para se conectar ao serviço externo e usar a conexão no código do seu agente. Consulte Conecte agentes a ferramentas de terceiros com os Serviços MCP.

O exemplo a seguir cria um recuperador que chama um índice vetorial hospedado fora do Databricks para um agente do tipo PyFunc.

  1. Crie uma Conexão do Unity Catalog para o serviço externo, neste caso, o Azure.

    SQL
    CREATE CONNECTION ${connection_name}
    TYPE HTTP
    OPTIONS (
    host 'https://example.search.windows.net',
    base_path '/',
    bearer_token secret ('<secret-scope>','<secret-key>')
    );
  2. Defina a ferramenta de retriever no código do agente usando a conexão do Unity Catalog. Este exemplo usa decoradores MLflow para habilitar o rastreamento do agente.

nota

Para estar em conformidade com o esquema do recuperador MLflow, a função do recuperador deve retornar um objeto List[Document] e usar o campo metadata na classe Document para adicionar atributos adicionais ao documento retornado, como doc_uri e similarity_score. Consulte o Documento do MLflow.

Python
import mlflow
import json

from mlflow.entities import Document
from typing import List, Dict, Any
from dataclasses import asdict

class VectorSearchRetriever:
"""
Class using Databricks AI 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.
"""
import requests
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
json = {
"count": true,
"select": "HotelId, HotelName, Description, Category",
"vectorQueries": [
{
"vector": query_vector,
"k": 7,
"fields": "DescriptionVector",
"kind": "vector",
"exhaustive": true,
}
],
}

response = requests.post(
f"{w.config.host}/api/2.0/unity-catalog/connections/{connection_name}/proxy/indexes/{self.azure_search_index}/docs/search?api-version=2023-07-01-Preview",
headers={
**w.config.authenticate(),
&quot;Content-Type&quot;: &quot;application/json&quot;,
},
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
  1. Para executar o retriever, execute o seguinte código Python.

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

Desenvolva um recuperador local

Desenvolva um retriever localmente usando a AI Bridge

Para criar uma ferramenta de recuperador Databricks AI Search localmente, use pacotes Databricks AI Bridge como databricks-langchain e databricks-openai. Esses pacotes incluem funções auxiliares como from_vector_search e from_uc_function para criar recuperadores a partir de recursos Databricks existentes.

Instale a versão mais recente de databricks-langchain que inclui o Databricks AI Bridge.

Bash
%pip install --upgrade databricks-langchain

O código a seguir cria o protótipo de uma ferramenta de recuperador que consulta um índice de pesquisa vetorial hipotético e o vincula a um LLM localmente para que você possa testar seu comportamento de chamada de ferramenta.

Forneça um tool_description descritivo para ajudar o agente a entender a ferramenta e determinar quando invocá-la.

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

Para cenários que usam índices de acesso direto ou índices Delta Sync usando embeddings autogerenciados, você deve configurar o VectorSearchRetrieverTool e especificar um modelo de embedding personalizado e uma coluna de texto. Consulte opções para fornecer embeddings.

O exemplo a seguir mostra como configurar um VectorSearchRetrieverTool com as keys columns e embedding.

Python
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={&quot;text_column LIKE&quot;: &quot;Databricks&quot;}, # 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.
)

Para obter detalhes adicionais, consulte a documentação da API para VectorSearchRetrieverTool.

Depois que sua ferramenta local estiver pronta, você poderá colocá-la diretamente em produção como parte do código do seu agente, ou migrá-la para uma função do Unity Catalog, que oferece melhor capacidade de descoberta e governança, mas apresenta certas limitações.

Query Databricks AI Search usando funções UC (obsoletas)

Query Databricks AI Search usando funções do UC (obsoleta)

nota

Databricks recommends MCP servers for most agent tools, but defining tools with Unity Catalog functions remains available for prototyping.

Você pode criar uma função do Unity Catalog que encapsula uma query de índice de pesquisa de AI. Esta abordagem:

  • Oferece suporte a casos de uso de produção com governança e possibilidade de descoberta.
  • Usa a função SQL vector_search() nos bastidores
  • Suporta rastreamento automático do MLflow
    • Você deve alinhar a saída da função ao esquema do recuperador MLflow usando os aliases page_content e metadata.
    • Quaisquer colunas de metadados adicionais devem ser adicionadas à coluna metadata usando a função de mapa SQL, em vez de como chaves de saída de nível superior.

Execute o código a seguir em um notebook ou editor SQL para criar a função:

SQL
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 AI Search index name here
index => 'catalog.schema.databricks_docs_index',
query => query,
num_results => 5
)

Para usar esta ferramenta de recuperador no seu agente de AI, envolva-a com UCFunctionToolkit. Isso permite o rastreamento automático através do MLflow, gerando automaticamente RETRIEVER tipos de span nos logs do MLflow.

Python
from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

As ferramentas de recuperação do Unity Catalog têm as seguintes ressalvas:

  • Clientes SQL podem limitar o número máximo de linhas ou bytes retornados. Para evitar o truncamento de dados, trunque os valores de coluna retornados pela UDF. Por exemplo, o senhor poderia usar substring(chunked_text, 0, 8192) para reduzir o tamanho de colunas de conteúdo grandes e evitar o truncamento de linhas durante a execução.
  • Como esta ferramenta é um invólucro para a função vector_search(), está sujeita às mesmas limitações que a função vector_search(). Consulte Limitações.

Para obter mais informações sobre UCFunctionToolkit, consulte a documentação do Unity Catalog.

Adicionar rastreamento a uma ferramenta de recuperação

Adicione rastreamento do MLflow para monitorar e depurar seu recuperador. O rastreamento permite visualizar entradas, saídas e metadados para cada passo de execução.

O exemplo anterior adiciona o decorador @mlflow.trace aos métodos __call__ e de análise. O decorador cria um span que começa quando a função é invocada e termina quando ela retorna. O MLflow registra automaticamente a entrada e a saída da função e quaisquer exceções geradas.

nota

Os usuários das bibliotecas LangChain, LlamaIndex e OpenAI podem usar o registro automático do MLflow, além de definir rastreamentos manualmente com o decorador. Consulte Adicionar rastreamentos a aplicativos: rastreamento automático e manual.

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

Para verificar se as aplicações downstream, como Agent Evaluation e AI Playground, renderizam o rastreamento do recuperador corretamente, certifique-se de que o decorador atende aos seguintes requisitos:

  • Use o esquema de span do MLflow retriever e verifique se a função retorna um objeto List[Document].
  • O nome do rastreamento e o nome retriever_schema devem corresponder para configurar o rastreamento corretamente. Consulte a seção a seguir para aprender a definir o esquema do recuperador.

Definir esquema do recuperador para verificar a compatibilidade com o MLflow

Se o rastreamento retornado do retriever ou span_type="RETRIEVER" não estiver em conformidade com o esquema de retriever padrão do MLflow, você deve mapear manualmente o esquema retornado para os campos esperados do MLflow. Isso verifica se o MLflow pode rastrear corretamente seu retriever e renderizar rastreamentos em aplicações downstream.

Para definir o esquema do retriever manualmente:

  1. Chame mlflow.models.set_retriever_schema quando você define seu agente. Use set_retriever_schema para mapear os nomes das colunas na tabela retornada para os campos esperados do MLflow, como primary_key, text_column e doc_uri.

    Python
    # 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. Especifique colunas adicionais no esquema do seu retriever fornecendo uma lista de nomes de colunas com o campo other_columns.

  3. Se houver vários recuperadores, é possível definir vários esquemas usando nomes exclusivos para cada esquema de recuperador.

O esquema do recuperador definido durante a criação do agente afeta aplicativos a jusante e fluxos de trabalho, como o aplicativo de revisão e os conjuntos de avaliação. Especificamente, a coluna doc_uri serve como o identificador principal para documentos retornados pelo recuperador.

Ler arquivos de um volume do Unity Catalog

Se seu agente precisar ler arquivos não estruturados (documentos de texto, relatórios, arquivos de configuração, etc.) armazenados em um volume do Unity Catalog, você pode criar ferramentas que usam a API de Arquivos do SDK do Databricks para listar e ler arquivos diretamente.

Os exemplos a seguir criam duas ferramentas que seu agente pode usar:

  • list_volume_files : Lista arquivos e diretórios no volume.
  • read_volume_file : Lê o conteúdo de um arquivo de texto do volume.

Instale a versão mais recente de databricks-langchain que inclui o Databricks AI Bridge.

Bash
%pip install --upgrade databricks-langchain
Python
from databricks.sdk import WorkspaceClient
from langchain_core.tools import tool

VOLUME = "<catalog>.<schema>.<volume>" # TODO: Replace with your volume
w = WorkspaceClient()


@tool
def list_volume_files(directory: str = "") -> str:
"""Lists files and directories in the Unity Catalog volume.
Provide a relative directory path, or leave empty to list the volume root."""
base = f"/Volumes/{VOLUME.replace('.', '/')}"
path = f"{base}/{directory.lstrip('/')}" if directory else base
entries = []
for f in w.files.list_directory_contents(path):
kind = "dir" if f.is_directory else "file"
size = f" ({f.file_size} bytes)" if not f.is_directory else ""
entries.append(f" [{kind}] {f.name}{size}")
return "\n".join(entries) if entries else "No files found."


@tool
def read_volume_file(file_path: str) -> str:
"""Reads a text file from the Unity Catalog volume.
Provide the path relative to the volume root, for example 'reports/q1_summary.txt'."""
base = f"/Volumes/{VOLUME.replace('.', '/')}"
full_path = f"{base}/{file_path.lstrip('/')}"
resp = w.files.download(full_path)
return resp.contents.read().decode("utf-8")

Vincule as ferramentas a um LLM e execute um loop de chamada de ferramenta:

Python
from databricks_langchain import ChatDatabricks
from langchain_core.messages import HumanMessage, ToolMessage

llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4-5")
llm_with_tools = llm.bind_tools([list_volume_files, read_volume_file])

messages = [HumanMessage(content="What files are in the volume? Can you read about_databricks.txt and summarize it in 2 sentences?")]
tool_map = {"list_volume_files": list_volume_files, "read_volume_file": read_volume_file}

for _ in range(5): # max iterations
response = llm_with_tools.invoke(messages)
messages.append(response)
if not response.tool_calls:
break
for tc in response.tool_calls:
result = tool_map[tc["name"]].invoke(tc["args"])
messages.append(ToolMessage(content=result, tool_call_id=tc["id"]))

print(response.content)