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

Prévia

Esse recurso está em Pré-lançamento público.

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 índices de pesquisa vetorial.

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

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

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

A função Unity Catalog databricks_docs_vector_search consulta um índice hipotético do Vector Search que contém a documentação da Databricks. Ele envolve a função Databricks SQL vector_search() e usa os aliases page_content e metadata para fazer a correspondência de sua saída com o esquemaMLflow retriever.

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 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 do site 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
  )

Essa ferramenta de recuperação tem as seguintes ressalvas:

• O MLflow rastreia essa 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 de agentes e o AI Playground, não mostrarão detalhes específicos do retriever, como links para chunks. Para obter mais informações sobre os 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, 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.

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

Recuperador de pesquisa vetorial com código de agente (PyFunc)

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

Este exemplo usa databricks-vectorsearch para criar um recuperador básico que executa uma pesquisa de similaridade de pesquisa vetorial com filtros. Ele usa decoradores do MLflow para permitir o rastreamento de agentes.

Observação

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.

Use o código a seguir no módulo do agente ou no Notebook 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 os 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 recuperadores sejam rastreados adequadamente e renderizados corretamente em aplicativos downstream, 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.

# 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 obter outras opções de configuração de métodos de rastreamento, consulte MLflow Tracing for agents.

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.

Observação

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 Usar o registro 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 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.

Filtrar resultados da pesquisa vetorial

O senhor pode limitar o escopo da pesquisa a um subconjunto de filtros de uso de dados do Vector Search.

O parâmetro filters em VectorSearchRetriever define as condições do filtro usando a especificação do filtro Databricks Vector 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 atendem às condições filters e score_threshold.

Próximos passos

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.