** Construir e rastrear ferramentas de recuperação para dados não estruturados

Use o Mosaic AI Agent Framework para criar ferramentas que permitem que os agentes de IA consultem dados não estruturados, como uma coleção de documentos. Esta página mostra como:

Desenvolver retrievers localmente
Criar um retriever usando as funções do Unity Catalog
Consultar índices vetoriais externos
Adicionar rastreamento MLflow para observabilidade

Para saber mais sobre as ferramentas do agente, consulte ferramentas de agentes de IA.

Desenvolva localmente ferramentas de busca vetorial com o AI Bridge

A maneira mais rápida de começar a construir uma ferramenta Databricks Vetor Search retriever é desenvolvê-la e testá-la localmente usando pacotes Databricks AI Bridge como databricks-langchain e databricks-openai.

LangChain/LangGraph

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

%pip install --upgrade databricks-langchain

O código a seguir prototipa uma ferramenta retriever 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 descritivo tool_description 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 representações autogeridas, deve configurar o VectorSearchRetrieverTool e deve especificar um modelo de representação personalizado e uma coluna de texto. Consulte as opções para fornecer embeddings.

O exemplo a seguir mostra como configurar um VectorSearchRetrieverTool com as chaves 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 os documentos da API para VectorSearchRetrieverTool.

OpenAI

Instale a versão mais recente do databricks-openai que inclui o Databricks AI Bridge.

%pip install --upgrade databricks-openai

O código a seguir prototipa um retriever que consulta um índice de pesquisa vetorial hipotético e o integra com os modelos GPT da OpenAI.

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

Para obter mais informações sobre recomendações do OpenAI para ferramentas, consulte documentação do OpenAI Function Calling.

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]
)

O exemplo a seguir mostra como configurar um VectorSearchRetrieverTool com as chaves 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 os documentos da API para VectorSearchRetrieverTool.

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

A seção a seguir mostra como migrar o retriever para uma função do Unity Catalog.

Ferramenta de pesquisa vetorial com funções do Unity Catalog

Você pode criar uma função Unity Catalog que encapsula uma consulta de índice do Mosaic AI Vector Search. Esta abordagem:

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

Execute o seguinte código em um bloco de anotações 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 Vector Search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

Para usar essa ferramenta retriever em seu agente de IA, envolva-a com UCFunctionToolkit. Isso permite o rastreamento automático através do MLflow, criando automaticamente 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 advertências:

Os clientes SQL podem limitar o número máximo de linhas ou bytes retornados. Para evitar o truncamento de dados, você deve truncar valores de coluna retornados pelo UDF. Por exemplo, você pode usar substring(chunked_text, 0, 8192) para reduzir o tamanho de colunas de conteúdo grandes e evitar o truncamento de linha durante a execução.
Uma vez que esta ferramenta é um wrapper para a função vector_search(), está sujeita às mesmas limitações que a função vector_search(). Consulte Limitações.

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

Retriever que consulta um índice vetorial hospedado fora do Databricks

Se seu índice de vetores estiver hospedado fora do Azure Databricks, você poderá criar uma conexão do Catálogo Unity para se conectar ao serviço externo e usar a conexão no código do agente. Veja Conexão de ferramentas do agente de IA a serviços externos.

O exemplo a seguir cria um retriever que acede a um índice vetorial armazenado fora da plataforma Databricks para um agente compatível com PyFunc.

Crie uma conexão de catálogo Unity com o serviço externo, neste caso, o Azure.

CREATE CONNECTION ${connection_name}
TYPE HTTP
OPTIONS (
  host 'https://example.search.windows.net',
  base_path '/',
  bearer_token secret ('<secret-scope>','<secret-key>')
);

Defina a ferramenta retriever no código do agente usando a conexão do Unity Catalog. Este exemplo usa decoradores MLflow para habilitar o rastreamento de agentes.

Observação

Para estar em conformidade com o esquema MLflow retriever, a função retriever deve retornar um List[Document] objeto e usar o metadata campo na classe Document para adicionar atributos adicionais ao documento retornado, como doc_uri e similarity_score. Consulte 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

Para executar o retriever, execute o seguinte código Python. Opcionalmente, você pode incluir filtros de Pesquisa Vetorial na solicitação para filtrar resultados.

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

Adicionar rastreamento a um retriever

Adicione o rastreamento MLflow para monitorizar e depurar o seu retriever. O rastreamento permite visualizar entradas, saídas e metadados para cada etapa da execução.

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

Observação

Os usuários das bibliotecas LangChain, LlamaIndex e OpenAI podem usar o registro automático MLflow, além de definir manualmente os traços com o decorador. Veja Adicionar traços às aplicações: rastreio 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 garantir que aplicações do processo subsequente, como a Avaliação de Agentes e o AI Playground, representem corretamente o traço do recuperador, certifique-se de que o decorador atenda aos seguintes requisitos:

Use (https://mlflow.org/docs/latest/tracing/tracing-schema.html#retriever-spans) e verifique se a função retorna um objeto List[Document].
O nome do rastro e o nome retriever_schema devem corresponder para configurar o rastro corretamente. Consulte a seção a seguir para saber como definir o esquema retriever.

Definir esquema retriever para garantir compatibilidade com MLflow

Se o rastreamento retornado do retriever ou span_type="RETRIEVER" não estiver em conformidade com o esquema retriever padrão do MLflow, você deverá mapear manualmente o esquema retornado para os campos esperados do MLflow. Isso garante que o MLflow possa rastrear adequadamente seu retriever e renderizar rastreamentos em aplicativos downstream.

Para definir o esquema retriever manualmente:

Chame mlflow.models.set_retriever_schema quando definir o 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_columne doc_uri.

# 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 retriever fornecendo uma lista de nomes de colunas com o other_columns campo.
Se você tiver vários retrievers, poderá definir vários esquemas usando nomes exclusivos para cada esquema retriever.

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

O aplicativo de revisão exibe o doc_uri para ajudar os revisores a avaliar as respostas e rastrear as origens do documento. Consulte Interface de Utilizador da App de Revisão.
Os conjuntos de avaliação são usados doc_uri para comparar os resultados do retriever com conjuntos de dados de avaliação predefinidos para determinar a recuperação e a precisão do retriever. Consulte Conjuntos de avaliação (MLflow 2).

Próximos passos

Depois de criar uma ferramenta de catálogo Unity, adicione-a ao seu agente. Consulte Criar uma ferramenta de agente.

Feedback

Esta página foi útil?

Last updated on 2025-12-09