Ferramentas de agente de IA de recuperação não estruturada

Importante

Esse recurso está em uma versão prévia.

Este artigo mostra como criar ferramentas de agente de IA para recuperação de dados não estruturados usando o Mosaic AI Agent Framework. Recuperadores não estruturados permitem que os agentes consultem fontes de dados não estruturadas, como um corpus de documento, usando índices de pesquisa de vetor.

Para saber mais sobre ferramentas para agentes de IA, confira ferramentas de agente de IA.

Desenvolver localmente as ferramentas recuperadoras 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 pacotes da Ponte de IA do Databricks como databricks-langchain e databricks-openai para adicionar rapidamente capacidades de recuperação de informações a um agente e experimentar com parâmetros de consulta. Essa abordagem permite iteração rápida durante o desenvolvimento inicial.

Depois que a ferramenta local estiver pronta, você poderá incorporá-la diretamente à produção como parte do código do agente ou migrá-la para uma função do Catálogo do Unity, que oferece melhor facilidade de descoberta e governança, mas tem certas limitações. Confira ferramenta de recuperação da Busca em Vetores com funções do Catálogo do Unity.

LangChain/LangGraph

O código a seguir cria um protótipo de uma ferramenta de recuperação e a vincula a um LLM localmente para que você possa conversar com o agente para testar o comportamento de chamada de ferramenta.

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

%pip install --upgrade databricks-langchain

Observação

Ao inicializar o VectorSearchRetrieverTool, os argumentos text_column e embedding são necessários para Índices de Sincronização Delta com inserções autogerenciadas e índices de acesso de vetor direto. Confira as opções para fornecer integrações.

from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool
vs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")

# 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?")

Para personalizar a chamada de ferramenta, passe argumentos adicionais para VectorSearchRetrieverTool:

from databricks_langchain import VectorSearchRetrieverTool

vs_tool = VectorSearchRetrieverTool(
  index_name, # Index name in the format 'catalog.schema.index'
  num_results, # Max number of documents to return
  columns, # List of columns to include in the search
  filters, # Filters to apply to the query
  query_type, # Query type ("ANN" or "HYBRID").
  tool_name, # Used by the LLM to understand the purpose of the tool
  tool_description, # Used by the LLM to understand the purpose of the tool
  text_column, # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
  embedding # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

OpenAI

O código a seguir prototipa uma ferramenta de pesquisa vetorial e a integra com os modelos GPT da OpenAI.

Para mais informações sobre as recomendações do OpenAI para ferramentas, consulte a documentação sobre Chamadas de Função do OpenAI .

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

%pip install --upgrade databricks-openai

Observação

Ao inicializar o VectorSearchRetrieverTool, os argumentos text_column e embedding são necessários para Índices de Sincronização Delta com incorporações autogerenciadas e Índices de Acesso Direto a Vetores. Confira as opções para fornecer integrações.

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

# Initialize OpenAI client
client = OpenAI(api_key=<your_API_key>)

# Call model with VectorSearchRetrieverTool defined
dbvs_tool = VectorSearchRetrieverTool(index_name="catalog.schema.my_index_name")
messages = [
  {"role": "system", "content": "You are a helpful assistant."},
  {
    "role": "user",
    "content": "Using the Databricks documentation, answer what is Spark?"
  }
]
first_response = client.chat.completions.create(
  model="gpt-4o",
  messages=messages,
  tools=[dbvs_tool.tool]
)

# Execute function code and parse the model's response and handle function calls.
tool_call = first_response.choices[0].message.tool_calls[0]
args = json.loads(tool_call.function.arguments)
result = dbvs_tool.execute(query=args["query"])  # For self-managed embeddings, optionally pass in openai_client=client

# Supply model with results – so it can incorporate them into its final response.
messages.append(first_response.choices[0].message)
messages.append({
  "role": "tool",
  "tool_call_id": tool_call.id,
  "content": json.dumps(result)
})
second_response = client.chat.completions.create(
  model="gpt-4o",
  messages=messages,
  tools=[dbvs_tool.tool]
)

Para personalizar a chamada de ferramenta, passe argumentos adicionais para VectorSearchRetrieverTool:

from databricks_openai import VectorSearchRetrieverTool

vs_tool = VectorSearchRetrieverTool(
    index_name, # Index name in the format 'catalog.schema.index'
    num_results, # Max number of documents to return
    columns, # List of columns to include in the search
    filters, # Filters to apply to the query
    query_type, # Query type ("ANN" or "HYBRID").
    tool_name, # Used by the LLM to understand the purpose of the tool
    tool_description, # Used by the LLM to understand the purpose of the tool
    text_column, # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
    embedding_model_name # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Ferramenta de recuperação de Busca em Vetores com funções do Catálogo do Unity

O exemplo a seguir cria uma ferramenta de recuperação de dados usando uma função do Catálogo do Unity para consultar dados de um indexador de Busca em vetores de IA do Mosaic.

A função Catálogo do Unity databricks_docs_vector_search consulta um índice hipotético de Busca em vetores que contém a documentação do Databricks. Essa função encapsula a função SQL do Databricks vector_search() e alinha sua saída com o esquema de recuperação MLflow. usando os aliases page_content e metadata.

Observação

Para estar em conformidade com o esquema de recuperação do MLflow, quaisquer 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 seguinte código em um notebook ou editor sql para criar a função:

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 de IA, encapsule-a com UCFunctionToolkit. Isso permite o rastreamento automático por meio do MLflow.

O rastreamento do MLflow captura informações detalhadas de execução para aplicativos de IA generativa. Ele registra entradas, saídas e metadados para cada etapa, ajudando você a depurar problemas e analisar o desempenho.

Ao usar UCFunctionToolkit, os recuperadores gerarão automaticamente tipos de intervalo RETRIEVER nos logs do MLflow, se a saída estiver em conformidade com o esquema de recuperador do MLflow. Confira Esquema de Rastreamento do MLflow.

Para obter mais informações sobre UCFunctionToolkit, confira Documentação do Catálogo do Unity.

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 truncamento de dados, você deve truncar os 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 truncamento de linha durante a execução.
• Como essa ferramenta é um wrapper para a função vector_search(), ela está sujeita às mesmas limitações que a função vector_search(). Confira Limitações.

Se este exemplo for inadequado para seu caso de uso, crie uma ferramenta de busca por vetor usando um código de agente personalizado.

Mecanismo de recuperação da busca em vetores com código de agente (PyFunc)

O exemplo a seguir cria um mecanismo de recuperação de busca em vetores para um agente no estilo PyFunc no código do agente.

Este exemplo usa [] databricks-vectorsearch [] para criar um recuperador básico que executa uma busca por similaridade de Vector Search [] com filtros []. Ele usa decoradores MLflow para habilitar o rastreamento de agente.

Observação

Para estar em conformidade com o esquema do recuperador do MLflow, a função de buscador deve retornar um tipo Documento e usar o campo metadata na classe Documento 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 recuperador, execute o seguinte código do Python. Opcionalmente, pode-se incluir filtros de Pesquisa de Vetor 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 mecanismo de recuperação

Para garantir que os buscadores sejam rastreados e renderizados corretamente em aplicativos downstream, chame mlflow.models.set_retriever_schema ao definir seu agente. Use set_retriever_schema para mapear os nomes de coluna 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 recuperador fornecendo uma lista de nomes de coluna com o campo other_columns.

Se você tiver vários buscadores, poderá definir vários esquemas usando nomes exclusivos para cada esquema de recuperação.

O esquema de recuperação definido durante a criação do agente afeta aplicativos downstream e fluxos de trabalho, como o aplicativo de revisão e os conjuntos de avaliação. Especificamente, a coluna doc_uri serve como o identificador primário para documentos retornados pelo buscador.

• O aplicativo de revisão exibe o doc_uri para ajudar os revisores a avaliar as respostas e rastrear as origens dos documentos. Confira a Interface do Usuário do Aplicativo de Análise.
• Conjuntos de avaliação usam doc_uri para comparar os resultados do mecanismo de recuperação com conjuntos de dados de avaliação predefinidos a fim de determinar o recall e a precisão do mecanismo de recuperação. Confira os Conjuntos de Avaliação.

Rastrear o mecanismo de recuperação

A funcionalidade de rastreamento do MLflow adiciona observabilidade ao capturar informações detalhadas sobre a execução do agente. Ele fornece um modo de registrar 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 a fim de criar um rastreamento para o mecanismo de recuperação e o analisador. Para mais opções de configuração de métodos de rastreamento, confira Rastreamento de MLflow para agentes.

O decorador cria um intervalo que começa quando a função é invocada e termina quando retorna. O MLflow registra automaticamente a entrada e a saída da função e todas as exceções geradas.

Observação

Os usuários das bibliotecas LangChain, LlamaIndex e OpenAI podem usar o registro automático do MLflow em vez de definir manualmente rastreamentos com o decorador. Confira 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 aplicativos posteriores, como a Avaliação do Agente e o Playground de IA, renderizem o rastreamento do mecanismo de recuperação corretamente, verifique se o decorador atende aos seguintes requisitos:

• Use span_type="RETRIEVER" e verifique se a função retorna o objeto List[Document]. Confira a abrangência do mecanismo de recuperação.
• O nome do rastreamento e o nome do retriever_schema devem ser iguais para configurar o rastreamento corretamente.

Filtrar resultados da busca em vetores

Você pode limitar o escopo de pesquisa a um subconjunto de dados usando um filtro de busca em vetores.

O parâmetro filters em VectorSearchRetriever define as condições de filtro usando a especificação de filtro da Busca em vetores do Databricks.

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

No 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 de filters e score_threshold.

Próximas etapas

Depois de criar uma ferramenta de agente de função do Catálogo do Unity, adicione a ferramenta a um agente de IA. Confira Adicionar ferramentas do Catálogo do Unity aos agentes.