Pular para o conteúdo principal

Recuperação não estruturada AI ferramentas de agente

Este artigo mostra como criar ferramentas de agente AI para recuperação de dados não estruturados usando o Mosaic AI Agent Framework. Os recuperadores não estruturados permitem que os agentes consultem fontes de dados não estruturadas, como um corpus de documentos, usando uma variedade de métodos de recuperação, incluindo a pesquisa baseada em vetores.

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 fácil de começar a desenvolver uma ferramenta de recuperação do Databricks Vector Search é localmente. Use o pacoteDatabricks AI Bridge como databricks-langchain e databricks-openai para adicionar rapidamente recursos de recuperação a um agente e experimentar parâmetros de consulta. Essa abordagem permite uma iteração rápida durante o desenvolvimento inicial.

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. Consulte a ferramenta de recuperação Vector Search com as funções do Unity Catalog.

nota

Para usar um índice vetorial externo hospedado fora da Databricks, consulte Vector Search retriever usando um índice vetorial hospedado fora da Databricks.

The following code prototypes a retriever tool and binds it to an LLM locally so you can chat with the agent to test its tool-calling behavior.

Install the latest version of databricks-langchain which includes Databricks AI Bridge.

Bash
%pip install --upgrade databricks-langchain

The following example queries a hypothetical vector search index that fetches content from Databricks product documentation.

Provide a clear and descriptive tool_description. The agent LLM uses the tool_description to understand the tool and determine when to invoke the tool.

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-meta-llama-3-1-70b-instruct")
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?")
nota

When initializing the VectorSearchRetrieverTool, the text_column and embedding arguments are required for Delta Sync Indexes with self-managed embeddings and Direct Vector Access Indexes. See options for providing embeddings.

For additional details, see the API docs for VectorSearchRetrieverTool

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

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

O exemplo a seguir cria uma ferramenta retriever usando uma função do Unity Catalog para consultar dados de um índice do Mosaic AI Vector Search.

A função databricks_docs_vector_search do Unity Catalog consulta um índice hipotético do Vector Search que contém a documentação da Databricks. Essa função envolve a função Databricks SQL vector_search() e alinha sua saída com o esquemaMLflow retriever. usando os aliases page_content e metadata.

nota

Para estar em conformidade com o esquema MLflow retriever, todas as colunas de metadados adicionais devem ser adicionadas à coluna metadata usando a função de mapaSQL, em vez de serem a 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 MLflow.

MLflow O rastreamento captura informações detalhadas sobre a execução de aplicativos do site AI. Ele logs entradas, saídas e metadados para cada etapa, ajudando o senhor a depurar problemas e analisar o desempenho.

Ao usar UCFunctionToolkit, os recuperadores geram automaticamente tipos de intervalo RETRIEVER em MLflow logs se a saída estiver em conformidade com o esquema do recuperador MLflow. Consulte Esquema de rastreamento do MLflow.

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

Python
from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

toolkit = UCFunctionToolkit(
    function_names=[
        "main.default.databricks_docs_vector_search"
    ]
)
tools = toolkit.tools

Essa ferramenta de recuperação tem 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.

Recuperação do Vector Search usando 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 o código do agente de conexão. Para obter mais informações, consulte Conectar as ferramentas do agente AI a um serviço externo.

O exemplo a seguir cria um recuperador Vector Search 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 que o senhor criou. Este exemplo usa decoradores do MLflow para habilitar o rastreamento de agentes.

nota

Para estar em conformidade com o esquema do retriever do MLflow, a função retriever deve retornar um tipo de documento e usar o campo metadata na classe Document para adicionar atributos adicionais ao documento retornado, como like doc_uri e similarity_score.

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)

Definir esquema do recuperador​

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

Para definir o esquema do recuperador manualmente, 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"],
)

Você também pode especificar colunas adicionais no esquema do seu retriever fornecendo uma lista de nomes de colunas com o campo other_columns.

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.

Rastreie o retriever​

MLflow O rastreamento adiciona observabilidade ao capturar informações detalhadas sobre a execução do agente. Ele fornece uma maneira de registrar as entradas, saídas e metadados associados a cada etapa intermediária de uma solicitação, permitindo que você identifique rapidamente a origem de bugs e comportamentos inesperados.

Este exemplo usa o decorador @mlflow .trace para criar um rastreamento para o recuperador e o analisador. Para conhecer outras opções de configuração de métodos de rastreamento, consulte Observabilidade do agente com o MLflow Tracing.

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, LlamaIndex e OpenAI biblioteca podem usar o registro automático do MLflow em vez de definir manualmente os traços com o decorador. Consulte Rastreamento automático.

Python
...
@mlflow.trace(span_type="RETRIEVER", name="vector_search")
def __call__(self, query: str) -> 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 o objeto List[Document]. Veja as extensões do Retriever.
  • O nome do rastreamento e o nome do retriever_schema devem coincidir para configurar o rastreamento corretamente.

Próximas etapas​

Depois que o senhor criar uma ferramenta de agente de função Unity Catalog, adicione a ferramenta a um agente AI. Consulte Adicionar ferramentas do Unity Catalog aos agentes.

Última atualização em