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.

For new Python agents, use AI Bridge (recommended) to query AI Search indexes directly. For existing MCP integrations, see the legacy workspace endpoint. You can also connect to external vector stores or read files from Unity Catalog volumes.

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)

Query um Índice de Pesquisa de AI usando MCP (legado)​

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

A URL de Workspace MCP legada para a 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 use este servidor, consulte Usar ferramentas MCP em um agente Python. Conceda ao agente SELECT no objeto protegível do Unity Catalog 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.