Pular para o conteúdo principal

Conectar agentes a dados não estruturados

Agentes geralmente precisam fazer query em dados não estruturados, como coleções de documentos, bases de conhecimento ou corpora de texto, para responder a perguntas e fornecer respostas com reconhecimento de contexto.

Para novos agentes Python, use AI Bridge (recomendado) para fazer query diretamente nos índices de AI Search. Para integrações MCP existentes, consulte o endpoint do workspace MCP. Você também pode se conectar a armazenamentos de vetores externos ou ler arquivos de volumes do Unity Catalog.

Desenvolva um retriever com o AI Bridge (recomendado)​

Para novos agentes Python que query índices do Databricks AI Search, o Databricks recomenda o AI Bridge. Use VectorSearchRetrieverTool de databricks-langchain ou databricks-openai para expor um índice como uma ferramenta de agente.

Antes de executar esses exemplos, faça login no seu workspace e escolha um Índice de Pesquisa de AI que você possa query. Substitua o nome do índice de exemplo pelo nome de três partes. Os primeiros exemplos usam embeddings gerenciados do Databricks e um endpoint de modelo compatível com chamadas de ferramentas.

Execução os comandos de instalação em um terminal. Em um notebook, use %pip em vez de pip e reinicie o Python após a instalação.

Instale a versão mais recente de databricks-langchain e LangGraph:

Bash
pip install --upgrade databricks-langchain langgraph

Crie uma ferramenta retriever e atribua a ela uma descrição que indique ao agente quando realizar a busca. Este exemplo testa a recuperação diretamente e, em seguida, execução um agente que pode chamar a ferramenta e usar seus resultados.

Python
from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks
from langgraph.prebuilt import create_react_agent

vs_tool = VectorSearchRetrieverTool(
index_name="catalog.schema.my_databricks_docs_index",
tool_name="databricks_docs_retriever",
tool_description="Search Databricks documentation to answer product questions.",
)
print(vs_tool.invoke("What is Databricks Agent Framework?"))

agent = create_react_agent(
ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
tools=[vs_tool],
)
result = agent.invoke({"messages": [{
"role": "user",
"content": "Search the documentation and explain Databricks Agent Framework.",
}]})
for message in result["messages"]:
print(message)

Verifique a conversa impressa para ver se há uma chamada de ferramenta, documentos recuperados e uma resposta baseada nesses documentos.

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={"text_column LIKE": "Databricks"}, # Filters to apply to the query
query_type="ANN", # Query type ("ANN" or "HYBRID").
tool_name="databricks_docs_retriever", # 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.

Após testar o recuperador, adicione-o ao seu agente personalizado.

Query um índice de pesquisa vetorial 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 Conectar-se a serviços HTTP externos.

The following example queries an Azure AI Search index through an HTTP connection. Você precisa de uma conexão com permissão para query esse índice e USE CONNECTION na conexão. Instale requests, databricks-sdk e mlflow antes de executar o exemplo em Python.

  1. Crie uma conexão do Unity Catalog com o Azure AI Search usando a configuração de conexão HTTP. Este exemplo usa um token do portador; escolha o método de autenticação exigido pela sua implantação do Azure AI Search.

    SQL
    CREATE CONNECTION azure_search_connection
    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 requests
from databricks.sdk import WorkspaceClient
from mlflow.entities import Document

class VectorSearchRetriever:
def __init__(self, connection_name, index_name):
self.workspace = WorkspaceClient()
self.connection_name = connection_name
self.index_name = index_name

@mlflow.trace(span_type="RETRIEVER", name="vector_search")
def __call__(self, query_vector: list[float], score_threshold: float = 0.0) -> list[Document]:
response = requests.post(
f"{self.workspace.config.host}/api/2.0/unity-catalog/connections/"
f"{self.connection_name}/proxy/indexes/{self.index_name}/docs/search",
params={"api-version": "2024-07-01"},
headers=self.workspace.config.authenticate(),
json={
"select": "HotelId,HotelName,Description,Category",
"vectorQueries": [{
"vector": query_vector,
"k": 7,
"fields": "DescriptionVector",
"kind": "vector",
"exhaustive": True,
}],
},
timeout=30,
)
response.raise_for_status()
return self.parse_results(response.json(), score_threshold)

@mlflow.trace(span_type="PARSER")
def parse_results(self, results, score_threshold) -> list[Document]:
return [
Document(
page_content=item.get("Description", ""),
id=item.get("HotelId"),
metadata={
"similarity_score": item.get("@search.score", 0),
"HotelName": item.get("HotelName"),
"Category": item.get("Category"),
},
)
for item in results.get("value", [])
if item.get("@search.score", 0) >= score_threshold
]
  1. Gere uma query embedding com o mesmo modelo e as mesmas dimensões usados pelo seu índice. Salve o array numérico como query-vector.json e, em seguida, faça a execução do retriever:

    Python
    import json
    from pathlib import Path

    query_vector = json.loads(Path("query-vector.json").read_text())
    retriever = VectorSearchRetriever("azure_search_connection", "hotels_vector_index")
    results = retriever(query_vector, score_threshold=0.1)
    for document in results:
    print(document.page_content)

Add tracing to a retriever tool​

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

LangChain, LlamaIndex, and OpenAI users da biblioteca podem usar o registro automático do MLflow, além de definir rastreamentos manualmente com o decorador. See Tracing overview.

Retorne uma lista de objetos Document com page_content e metadata, como o exemplo acima faz. Isso segue o esquema de intervalo do recuperador do MLflow, para que as ferramentas de avaliação e revisão possam reconhecer os documentos recuperados.

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.

Escolha um volume existente e conceda à identidade do seu agente USE CATALOG, USE SCHEMA e READ VOLUME. Substitua <catalog>.<schema>.<volume> no código. Para a tarefa de amostra, use um arquivo de texto UTF-8 chamado about_databricks.txt ou altere o prompt para um arquivo em seu volume.

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 langgraph
Python
from pathlib import PurePosixPath
from databricks.sdk import WorkspaceClient
from langchain_core.tools import tool

VOLUME = "<catalog>.<schema>.<volume>"
w = WorkspaceClient()


def volume_path(relative_path: str = "") -> str:
path = PurePosixPath(relative_path)
if path.is_absolute() or ".." in path.parts:
raise ValueError("Use a path relative to the configured volume.")
return str(PurePosixPath("/Volumes", *VOLUME.split(".")) / path)


@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."""
path = volume_path(directory)
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'."""
full_path = volume_path(file_path)
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"]))

else:
raise RuntimeError("The agent reached the tool-calling limit.")

print(response.content)

Faça uma query em um índice de Pesquisa de AI usando MCP​

O servidor AI Search MCP usa um endpoint MCP de workspace. Ele requer um índice com incorporações gerenciadas do Databricks. Consulte Criar Endpoint e índices de Pesquisa de AI.

O URL MCP do workspace para AI Search é 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 criar e implantar um agente que usa este servidor, consulte Chamar ferramentas MCP de um agente Python. Conceda ao agente SELECT no item do Unity Catalog que pode ser protegido do índice.

A Pesquisa de AI era anteriormente a Pesquisa Vetorial. O prefixo de URL /api/2.0/mcp/vector-search/ anterior ainda funciona.

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

nota

Para novos agentes Python, use o AI Bridge. O exemplo a seguir é uma referência para integrações existentes.

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
)

Siga a configuração da ferramenta de função e, em seguida, envolva este recuperador com UCFunctionToolkit para usá-lo em seu agente. Isso habilita o rastreamento automático por meio do MLflow gerando automaticamente tipos de intervalo RETRIEVER nos Logs do MLflow.

Python
from databricks_langchain import UCFunctionToolkit
from unitycatalog.ai.core.databricks import DatabricksFunctionClient

toolkit = UCFunctionToolkit(
function_names=[
"main.default.databricks_docs_vector_search"
],
client=DatabricksFunctionClient(),
)
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.

Esquemas de retriever personalizados (legados)​

mlflow.models.set_retriever_schema foi descontinuado. Para novos recuperadores, retorne Document objetos usando o esquema de recuperador padrão. A configuração a seguir é para integrações existentes que retornam campos personalizados.

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
    import mlflow

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