Crie uma ferramenta de recuperação de busca vetorial

Prévia

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

Saiba como usar o Mosaic AI Agent Framework para criar recuperadores. Um recuperador é um tipo de ferramenta de agente que localiza e retorna documentos relevantes usando um índice de pesquisa vetorial. Os recuperadores são um componente central dos aplicativos RAG (Retrieval Augmented Generation).

Requisitos

Exemplo de recuperador PyFunc

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

A função de recuperação deve retornar um tipo de documento. Use 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 corretamente, chame mlflow.models.set_retriever_schema quando você define seu agente em código. 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"],
)

Observação

A coluna doc_uri é especialmente importante ao avaliar o desempenho do retriever. doc_uri é o identificador principal dos documentos retornados pelo recuperador, permitindo compará-los com os conjuntos de avaliação da verdade básica. Consulte Conjuntos de avaliação.

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.

Rastreie o retriever

MLflow O rastreamento acrescenta observabilidade ao capturar informações detalhadas sobre a execução do agente. Ele oferece uma maneira de registrar as entradas, saídas e metadados associados a cada passo intermediário de uma solicitação, o que permite identificar facilmente 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 geradas por ela.

Observação

LangChainOs usuários do MLflow, do LlamaIndex e do OpenAI biblioteca podem usar o autologging 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, renderizem 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 de filtro usando a especificação de 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.

Exemplos de aplicativos do Retriever

Consulte o repositório GitHub do genai-cookbook para obter exemplos do agente AI que usam retrievers:

Próximos passos