Construir e monitorar ferramentas de recuperação para dados não estruturados

Use o Mosaic AI Agent Framework para criar ferramentas que permitem aos agentes de IA consultar dados não estruturados, como uma coleção de documentos. Esta página mostra como:

• Desenvolver recuperadores localmente
• Criar um recuperador usando funções do Catálogo do Unity
• Consultar índices de vetores externos
• Adicionar rastreamento de MLflow para observabilidade

Para saber mais sobre as ferramentas do agente, confira Ferramentas de agente de IA.

Desenvolver localmente as ferramentas de busca do Vector Search com a AI Bridge

A maneira mais rápida de começar a criar uma ferramenta recuperadora de busca em vetores do Databricks é desenvolvê-la e testá-la localmente usando pacotes da Ponte de IA do Databricks como databricks-langchain e databricks-openai.

LangChain/LangGraph

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

%pip install --upgrade databricks-langchain

O código a seguir é um protótipo de uma ferramenta de busca que consulta um índice de pesquisa de vetor hipotético e o conecta a um LLM localmente, para que você possa testar seu comportamento de acionamento da ferramenta.

Forneça um descritivo tool_description para ajudar o agente a entender a ferramenta e determinar quando invocá-la.

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-claude-3-7-sonnet")
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 cenários que usam índices de acesso direto ou índices de Sincronização Delta usando embeddings autogeridos, você deve configurar o VectorSearchRetrieverTool, especificar um modelo de embedding personalizado e uma coluna de texto. Confira as opções para fornecer integrações.

O exemplo a seguir mostra como configurar um VectorSearchRetrieverTool com chaves columns e embedding.

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

Para obter detalhes adicionais, consulte os documentos da API para VectorSearchRetrieverTool.

OpenAI

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

%pip install --upgrade databricks-openai

O código a seguir protótipa um recuperador que consulta um índice de pesquisa de vetor hipotético e o integra aos modelos GPT do OpenAI.

Forneça um descritivo tool_description para ajudar o agente a entender a ferramenta e determinar quando invocá-la.

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

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

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

# Initialize the retriever tool
dbvs_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"
)

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 cenários que usam índices de acesso direto ou índices de Sincronização Delta usando embeddings autogeridos, você deve configurar o VectorSearchRetrieverTool, especificar um modelo de embedding personalizado e uma coluna de texto. Confira as opções para fornecer integrações.

O exemplo a seguir mostra como configurar um VectorSearchRetrieverTool com chaves columns e embedding.

from databricks_openai import VectorSearchRetrieverTool

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_model_name="databricks-bge-large-en" # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

Para obter detalhes adicionais, consulte os documentos da API para VectorSearchRetrieverTool.

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.

A seção a seguir mostra como migrar o recuperador para uma função do Catálogo do Unity.

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

Você pode criar uma função de Catálogo do Unity que encapsula uma consulta de Índice de busca em vetores da IA do Mosaic. Essa abordagem:

• Dá suporte a casos de uso de produção com governança e capacidade de descoberta
• Usa a função SQL vector_search() nos bastidores
• Dá suporte ao rastreamento automático do MLflow
  • Você deve alinhar a saída da função ao Esquema de recuperação do MLflow usando os aliases page_content e metadata.
  • Quaisquer colunas de metadados adicionais devem ser adicionadas à metadata coluna 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 no seu agente de IA, encapsule-a com UCFunctionToolkit. Isso permite o rastreamento automático por meio do MLflow ao gerar automaticamente tipos de span RETRIEVER nos logs do MLflow..

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

As ferramentas de recuperação do Catálogo do Unity têm 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.

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

Recuperador que consulta um índice de vetor hospedado fora do Databricks

Se o índice de vetor estiver hospedado fora do Azure Databricks, você poderá criar uma conexão do Unity Catalog para se conectar ao serviço externo e usar a conexão no código do agente. Consulte Conectar ferramentas de agente de IA a serviços externos.

O exemplo a seguir cria um recuperador que chama um índice de vetor hospedado fora do Databricks para um agente compatível com PyFunc.

1. Crie uma Conexão de Catálogo do Unity com o serviço externo, nesse caso, o Azure.

   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 de recuperação no código do agente usando a conexão do Unity Catalog. Este exemplo usa decoradores MLflow para habilitar o rastreamento de agente.

   Observação

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

   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
   

3. 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 = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
   results = retriever(query, score_threshold=0.1)
   

Adicionar rastreamento a um recuperador

Adicione o rastreamento do MLflow para monitorar e depurar o recuperador. O rastreamento permite exibir entradas, saídas e metadados para cada etapa de execução.

O exemplo anterior adiciona o decorador @mlflow.trace aos métodos __call__ e de análise. 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, além de definir manualmente os traços com o decorador. Consulte Instrumente seu aplicativo: abordagens de rastreamento.

import mlflow
from mlflow.entities import Document

## This code snippet has been truncated for brevity, see the full retriever example above
class VectorSearchRetriever:
  ...

  # Create a RETRIEVER span. The span name must match the retriever schema name.
  @mlflow.trace(span_type="RETRIEVER", name="vector_search")
  def __call__(...) -> List[Document]:
    ...

  # Create a PARSER span.
  @mlflow.trace(span_type="PARSER")
  def parse_results(...) -> List[Document]:
    ...

Para garantir que aplicativos downstream, como Avaliação de Agentes e o Playground de IA, renderizem corretamente o rastreamento do recuperador, verifique se o decorador atende aos seguintes requisitos:

• Use (https://mlflow.org/docs/latest/tracing/tracing-schema.html#retriever-spans) e verifique se a função retorna um objeto List[Document].
• O nome do rastreamento e o nome retriever_schema devem corresponder para configurar o rastreamento corretamente. Consulte a próxima seção para aprender a definir o esquema do recuperador.

Definir o esquema do retriever para garantir a compatibilidade do MLflow

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

Para definir o esquema do recuperador manualmente:

1. Chame mlflow.models.set_retriever_schema quando 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_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"],
   )
   

2. Especifique colunas adicionais no esquema do seu recuperador fornecendo uma lista de nomes de colunas usando o campo other_columns.

3. 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. Consulte Interface do Usuário do Aplicativo de Análise.
• Os conjuntos de avaliação usam doc_uri para comparar os resultados do sistema de recuperação com conjuntos de dados de avaliação predefinidos, a fim de determinar o recall e a precisão do sistema de recuperação. Consulte Conjuntos de Avaliação (MLflow 2).

Próximas etapas

Depois de construir o recuperador, a etapa final é integrá-lo na definição de um agente de IA. Saiba como adicionar uma ferramenta a um agente, consulte Adicionar ferramentas do Catálogo do Unity aos agentes.