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:
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.
-
Crie uma Conexão do Unity Catalog para o serviço externo, neste caso, o Azure.
SQLCREATE CONNECTION ${connection_name}
TYPE HTTP
OPTIONS (
host 'https://example.search.windows.net',
base_path '/',
bearer_token secret ('<secret-scope>','<secret-key>')
); -
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.
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.
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(),
"Content-Type": "application/json",
},
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
-
Para executar o retriever, execute o seguinte código Python.
Pythonretriever = 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.
- LangChain/LangGraph
- OpenAI
Instale a versão mais recente de databricks-langchain que inclui o Databricks AI Bridge.
%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.
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.
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.
)
Para obter detalhes adicionais, consulte a documentação da API para VectorSearchRetrieverTool.
Instale a versão mais recente de databricks-openai que inclui o Databricks AI Bridge.
%pip install --upgrade databricks-openai
O código a seguir prototipa um recuperador que consulta um índice hipotético de busca vetorial e o integra aos modelos GPT da OpenAI.
Forneça um tool_description descritivo para ajudar o agente a entender a ferramenta e determinar quando invocá-la.
Para obter mais informações sobre as recomendações do OpenAI para ferramentas, consulte a documentação de Chamada de Função do 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]
)
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.
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.
)
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)
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_contentemetadata. - Quaisquer colunas de metadados adicionais devem ser adicionadas à coluna
metadatausando a função de mapa SQL, em vez de como chaves de saída de nível superior.
- Você deve alinhar a saída da função ao esquema do recuperador MLflow usando os aliases
Execute o código a seguir em um notebook ou editor SQL para criar a função:
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.
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çãovector_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.
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.
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_schemadevem 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:
-
Chame mlflow.models.set_retriever_schema quando você define seu agente. Use
set_retriever_schemapara mapear os nomes das colunas na tabela retornada para os campos esperados do MLflow, comoprimary_key,text_columnedoc_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"],
) -
Especifique colunas adicionais no esquema do seu retriever fornecendo uma lista de nomes de colunas com o campo
other_columns. -
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.
- O aplicativo de revisão exibe
doc_uripara ajudar os revisores a avaliar as respostas e rastrear as origens dos documentos. Consulte Interface do Usuário do Aplicativo de Revisão. - Conjuntos de avaliação usam
doc_uripara comparar resultados do recuperador com datasets de avaliação predefinidos para determinar o recall e a precisão do recuperador. Consulte conjuntos de avaliação (MLflow 2).
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.
- LangChain/LangGraph
- OpenAI
Instale a versão mais recente de databricks-langchain que inclui o Databricks AI Bridge.
%pip install --upgrade databricks-langchain
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:
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)
Instale a versão mais recente de databricks-openai que inclui o Databricks AI Bridge.
%pip install --upgrade databricks-openai
from databricks.sdk import WorkspaceClient
from databricks_openai import DatabricksOpenAI
import json
VOLUME = "<catalog>.<schema>.<volume>" # TODO: Replace with your volume
w = WorkspaceClient()
client = DatabricksOpenAI()
# Define the tool specifications
tools = [
{
"type": "function",
"function": {
"name": "list_volume_files",
"description": "Lists files and directories in the Unity Catalog volume. Provide a relative directory path, or leave empty to list the volume root.",
"parameters": {
"type": "object",
"properties": {
"directory": {
"type": "string",
"description": "Relative directory path within the volume. Leave empty for root.",
}
},
"required": [],
},
},
},
{
"type": "function",
"function": {
"name": "read_volume_file",
"description": "Reads a text file from the Unity Catalog volume. Provide the path relative to the volume root, for example 'reports/q4_summary.txt'.",
"parameters": {
"type": "object",
"properties": {
"file_path": {
"type": "string",
"description": "Path to the file relative to the volume root.",
}
},
"required": ["file_path"],
},
},
},
]
def execute_tool(name: str, args: dict) -> str:
base = f"/Volumes/{VOLUME.replace('.', '/')}"
if name == "list_volume_files":
directory = args.get("directory", "")
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."
elif name == "read_volume_file":
full_path = f"{base}/{args['file_path'].lstrip('/')}"
resp = w.files.download(full_path)
return resp.contents.read().decode("utf-8")
return f"Unknown tool: {name}"
# Call the model with tools
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "List the files in the volume, then read about_databricks.txt and summarize it."},
]
response = client.chat.completions.create(
model="databricks-claude-sonnet-4-5", messages=messages, tools=tools
)
# Execute tool calls and send results back
while response.choices[0].finish_reason == "tool_calls":
messages.append(response.choices[0].message)
for tool_call in response.choices[0].message.tool_calls:
args = json.loads(tool_call.function.arguments)
result = execute_tool(tool_call.function.name, args)
messages.append(
{"role": "tool", "tool_call_id": tool_call.id, "content": result}
)
response = client.chat.completions.create(
model="databricks-claude-sonnet-4-5", messages=messages, tools=tools
)
print(response.choices[0].message.content)