Pular para o conteúdo principal

Crie e rastreie ferramentas de recuperação para dados não estruturados

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

Para saber mais sobre as ferramentas do agente, consulte AI agent tools.

Desenvolver localmente as ferramentas de recuperação do Vector Search com o AI Bridge​

A maneira mais rápida de começar a criar uma ferramenta de recuperação do Databricks Vector Search é desenvolvê-la e testá-la localmente usando Databricks AI Bridge pacote como databricks-langchain e databricks-openai.

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

Bash
%pip install --upgrade databricks-langchain

O código a seguir cria um protótipo de uma ferramenta de recuperação que consulta um índice de pesquisa vetorial hipotético e o vincula a um LLM localmente para que o senhor possa testar o comportamento de chamada da 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-3-7-sonnet")
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 do Delta Sync usando incorporações autogerenciadas, o senhor deve configurar o VectorSearchRetrieverTool e especificar um modelo de incorporação personalizado e uma coluna de texto. Veja as opções para fornecer incorporações.

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

Quando sua ferramenta local estiver pronta, o senhor pode produzi-la diretamente como parte do código do agente ou migrá-la para uma função do 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 recuperador para uma função do Unity Catalog.

Ferramenta de recuperação do Vector Search com funções do Unity Catalog​

O senhor pode criar uma função do Unity Catalog que envolva uma consulta de índice do Mosaic AI Vector Search. Essa abordagem:

  • Oferece suporte a casos de uso de produção com governança e visibilidade
  • Usa a função SQL vector_search() por trás do capô
  • Oferece suporte ao rastreamento automático de MLflow
    • O senhor deve alinhar a saída da função ao esquema de recuperação do MLflow usando os aliases page_content e metadata.
    • Quaisquer colunas de metadados adicionais devem ser adicionadas à coluna metadata usando a função de mapaSQL, e não como chave de saída de nível superior.

Execute o código a seguir em um editor do Notebook ou 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 Vector Search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

Para usar essa ferramenta de recuperação em seu agente AI, envolva-a com UCFunctionToolkit. Isso permite o rastreamento automático por meio do site MLflow, gerando automaticamente os tipos de extensão RETRIEVER em MLflow logs.

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:

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

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

Retriever que consulta um índice vetorial hospedado fora da Databricks​

Se o seu índice vetorial estiver hospedado fora da Databricks, o senhor poderá criar uma conexão do Unity Catalog para se conectar ao serviço externo e usar a conexão no código do agente. Consulte Conectar as ferramentas do agente AI ao serviço externo.

O exemplo a seguir cria um retriever que chama um índice de vetor hospedado fora do Databricks para um agente com sabor de PyFunc.

  1. Crie uma conexão do Unity Catalog com 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 retriever no código do agente usando a conexão do Unity Catalog. Este exemplo usa decoradores do MLflow para habilitar o rastreamento de agentes.

nota

Para estar em conformidade com o esquema retriever do MLflow, a função retriever 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 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 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

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

    Python
    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 do MLflow para monitorar e depurar seu recuperador. O rastreamento permite que o senhor view entradas, saídas e metadados para cada etapa da execução.

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

nota

LangChainOs usuários do OpenAI biblioteca, do LlamaIndex e do MLflow podem usar o registro automático, além de definir manualmente os traços com o decorador. Consulte Adicionar MLflow Tracing aos agentes AI.

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 garantir que os aplicativos downstream, como o Agent Evaluation e o AI Playground, processem o rastreamento do retriever corretamente, certifique-se de que o decorador atenda aos seguintes requisitos:

  • Use span_type="RETRIEVER" e certifique-se de que a função retorne um objeto List [Document]. Veja as extensões do Retriever.
  • O nome do rastreamento e o nome retriever_schema devem coincidir para configurar o rastreamento corretamente. Consulte a seção a seguir para saber como definir o esquema do retriever.

Definir o esquema do retriever para garantir a compatibilidade com o MLflow​

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

Para definir o esquema do recuperador 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 você tiver vários recuperadores, poderá definir vários esquemas usando nomes exclusivos para cada esquema de recuperação.

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

  • O aplicativo de avaliação exibe o doc_uri para ajudar os revisores a avaliar as respostas e rastrear as origens dos documentos. Consulte Revisar a interface do usuário do aplicativo.
  • Os conjuntos de avaliação usam o site doc_uri para comparar os resultados do recuperador com o conjunto de dados de avaliação predefinido para determinar a recuperação e a precisão do recuperador. Consulte Conjuntos de avaliação.

Próximas etapas​

Depois de criar seu retriever, a etapa final é integrá-lo a uma definição de agente AI. Para saber como adicionar uma ferramenta a um agente, consulte Adicionar ferramentas do Unity Catalog a agentes.

Última atualização em