Compartilhar via

Ferramentas de agentes de IA para recuperação não estruturada

  • Artigo

Importante

Este recurso está em Public Preview.

Este artigo mostra como criar ferramentas de agente de IA 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 índices de pesquisa vetorial.

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

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

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

A função Unity Catalog databricks_docs_vector_search consulta um índice hipotético de Pesquisa Vetorial contendo documentação do Databricks. Ele encapsula a função Databricks SQL vetor_search() e usa os aliases page_content e metadata para corresponder sua saída ao esquema de recuperação MLflow .

Observação

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 mapa SQL , em vez de como chaves de saída de nível superior.

Execute o código a seguir em um bloco de anotações ou editor 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
  )

Esta ferramenta de recuperação tem as seguintes limitações:

  • MLflow rastreia esta função do Unity Catalog como um tipo de intervalo TOOL em vez de um tipo de intervalo RETRIEVER. Como resultado, os aplicativos downstream do Agent Framework, como o aplicativo de revisão do agente e o AI Playground, não mostrarão detalhes específicos do retriever, como links para blocos. Para obter mais informações sobre tipos de span, consulte MLflow Tracing Schema.
  • 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.

Se este exemplo não for adequado para o seu caso de uso, crie uma ferramenta de pesquisa vetorial usando código de agente personalizado.

Buscador Vetorial com código de agente (PyFunc)

O exemplo a seguir cria um recuperador de pesquisa vetorial para um agente no estilo PyFunc dentro do código do agente.

Este exemplo usa databricks-vectorsearch para criar um recuperador básico que executa uma pesquisa de semelhança de Pesquisa Vetorial com filtros. Ele usa decoradores MLflow para habilitar o rastreamento de agentes.

Observação

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

Use o código a seguir no módulo do agente ou no bloco de anotações do agente.

import mlflow
import json

from mlflow.entities import Document
from typing import List, Dict, Any
from dataclasses import asdict
from databricks.vector_search.client import VectorSearchClient

class VectorSearchRetriever:
    """
    Class using Databricks Vector Search to retrieve relevant documents.
    """
    def __init__(self):
        self.vector_search_client = VectorSearchClient(disable_notice=True)
        # TODO: Replace this with the list of column names to return in the result when querying Vector Search
        self.columns = ["chunk_id", "text_column", "doc_uri"]
        self.vector_search_index = self.vector_search_client.get_index(
            index_name="catalog.schema.chunked_docs_index"
        )
        mlflow.models.set_retriever_schema(
            name="vector_search",
            primary_key="chunk_id",
            text_column="text_column",
            doc_uri="doc_uri"
        )

    @mlflow.trace(span_type="RETRIEVER", name="vector_search")
    def __call__(
        self,
        query: str,
        filters: Dict[Any, Any] = None,
        score_threshold = None
    ) -> List[Document]:
        """
        Performs vector search to retrieve relevant chunks.
        Args:
            query: Search query.
            filters: Optional filters to apply to the search. Filters must follow the Databricks Vector Search filter spec
            score_threshold: Score threshold to use for the query.

        Returns:
            List of retrieved Documents.
        """

        results = self.vector_search_index.similarity_search(
            query_text=query,
            columns=self.columns,
            filters=filters,
            num_results=5,
            query_type="ann"
        )

        documents = self.convert_vector_search_to_documents(
            results, 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 = []
        column_names = [column["name"] for column in vs_results.get("manifest", {}).get("columns", [])]
        result_row_count = vs_results.get("result", {}).get("row_count", 0)

        if result_row_count > 0:
            for item in vs_results["result"]["data_array"]:
                metadata = {}
                score = item[-1]

                if score >= score_threshold:
                    metadata["similarity_score"] = score
                    for i, field in enumerate(item[:-1]):
                        metadata[column_names[i]] = field

                    page_content = metadata.pop("text_column", None)

                    if page_content:
                        doc = Document(
                            page_content=page_content,
                            metadata=metadata
                        )
                        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 = "What is Databricks?"
filters={"text_column LIKE": "Databricks"},
results = retriever(query, filters=filters, score_threshold=0.1)

Definir esquema do recuperador

Para garantir que os retrievers sejam rastreados corretamente e renderizados corretamente em aplicativos downstream, chame mlflow.models.set_retriever_schema quando definir 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"],
)

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

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 Review Appinterface de utilizador da aplicação.
  • Os conjuntos de avaliação usam doc_uri para comparar os resultados do retriever com conjuntos de dados de avaliação predefinidos para determinar a revocação e a precisão do retriever. Consulte Conjuntos de avaliação.

Rastrear o retriever

O rastreamento MLflow adiciona observabilidade capturando informações detalhadas sobre a execução do seu 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 a origem de bugs e comportamentos inesperados rapidamente.

Este exemplo usa o decorador @mlflow.trace para criar um rastreamento para o retriever e o analisador. Para obter outras opções para configurar métodos de rastreamento, consulte MLflow Tracing for agents.

O decorador cria um de 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 utilizadores das bibliotecas LangChain, LlamaIndex e OpenAI podem usar o registo automático do MLflow em vez de definir manualmente os traces com o decorador. Veja Use o registo automático para adicionar rastreamentos aos seus agentes.

...
@mlflow.trace(span_type="RETRIEVER", name="vector_search")
def __call__(self, query: str) -> List[Document]:
  ...

Para garantir que aplicativos de fluxo descendente, como a Avaliação de Agentes e o AI Playground, executem corretamente o rastreamento do recuperador, certifique-se de que o decorador atenda aos seguintes requisitos:

  • Use span_type="RETRIEVER" e verifique se a função retorna List[Document] objeto. Veja os períodos do Retriever .
  • O nome do rastreamento e o nome do retriever_schema devem corresponder para configurar o rastreamento corretamente.

Filtrar resultados da Pesquisa Vetorial

Você pode limitar o escopo da pesquisa a um subconjunto de dados usando um filtro de Pesquisa Vetorial.

O parâmetro filters no VectorSearchRetriever define as condições do filtro usando a especificação de filtro Databricks Vetor Search.

filters = {"text_column LIKE": "Databricks"}

Dentro do método __call__, o dicionário de filtros é passado diretamente para a função similarity_search:

results = self.vector_search_index.similarity_search(
    query_text=query,
    columns=self.columns,
    filters=filters,
    num_results=5,
    query_type="ann"
)

Após a filtragem inicial, o parâmetro score_threshold fornece filtragem adicional definindo uma pontuação mínima de similaridade.

if score >= score_threshold:
    metadata["similarity_score"] = score

O resultado final inclui documentos que cumprem as condições filters e score_threshold.

Próximos passos

Depois de criar uma ferramenta de agente de função no Unity Catalog, adicione a ferramenta a um agente de IA. Consulte para adicionar ferramentas do Catálogo Unity aos agentes.