Compartir vía

Herramientas de agente de IA de recuperación no estructurada

  • Artículo

Importante

Esta característica está en versión preliminar pública.

En este artículo se muestra cómo crear herramientas de agente de IA para la recuperación de datos no estructurados mediante el marco del agente de IA de Mosaico. Los recuperadores no estructurados permiten a los agentes consultar orígenes de datos no estructurados, como un corpus de documentos, mediante índices de búsqueda vectorial.

Para obtener más información sobre las herramientas del agente, consulte Herramientas del agente de IA.

Desarrollo local de herramientas de recuperación de búsqueda vectorial con AI Bridge

La manera más fácil de empezar a desarrollar una herramienta de recuperación de Búsqueda vectorial de Databricks es localmente. Utilice paquetes de Databricks AI Bridge como databricks-langchain y databricks-openai para agregar rápidamente capacidades de recuperación a un agente y experimentar con parámetros de consulta. Este enfoque permite una iteración rápida durante el desarrollo inicial.

Una vez que la herramienta local esté lista, puede producirla directamente como parte del código del agente o migrarla a una función de catálogo de Unity, lo que proporciona una mejor detectabilidad y gobernanza, pero tiene ciertas limitaciones. Consulte la Herramienta de vector de búsqueda con funciones de Unity Catalog.

LangChain/LangGraph

El código siguiente prototipa una herramienta de recuperación y la vincula a un LLM localmente para que pueda chatear con el agente y probar el comportamiento del agente al invocar herramientas.

Instale la versión más reciente de databricks-langchain que incluye Databricks AI Bridge.

%pip install --upgrade databricks-langchain

Nota:

Al inicializar el VectorSearchRetrieverTool, los argumentos text_column y embedding son necesarios para los índices de sincronización delta con incrustaciones autoadministradas y índices de acceso vectorial directo. Consulte las opciones para proporcionar incrustaciones.

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 las llamadas a herramientas, pase argumentos adicionales a 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

El código siguiente crea prototipos de una herramienta de recuperación de búsqueda vectorial y la integra con los modelos GPT de OpenAI.

Para obtener más información sobre las recomendaciones de OpenAI para herramientas, consulte Documentación sobre llamadas a funciones de OpenAI.

Instale la versión más reciente de databricks-openai que incluye Databricks AI Bridge.

%pip install --upgrade databricks-openai

Nota:

Al inicializar el VectorSearchRetrieverTool, los argumentos text_column y embedding son necesarios para los índices de sincronización delta con incrustaciones autoadministradas y índices de acceso vectorial directo. Consulte las opciones para proporcionar incrustaciones.

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 las llamadas a herramientas, pase argumentos adicionales a 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.
)

Herramienta del recuperador de vector de búsqueda con funciones de Unity Catalog

En el ejemplo siguiente se crea una herramienta de búsqueda mediante una función de Unity Catalog para consultar datos de un índice de búsqueda vectorial de Mosaic AI .

La función Catálogo de Unity databricks_docs_vector_search consulta un índice hipotético de búsqueda vectorial que contiene la documentación de Databricks. Esta función envuelve la función SQL de Databricks vector_search() y alinea su salida con el esquema del recuperador de MLflow. mediante los alias page_content y metadata.

Nota:

Para ajustarse al esquema del recuperador de MLflow, se deben agregar columnas de metadatos adicionales a la columna metadata mediante la función de asignación de SQL, en lugar de como claves de salida de nivel superior.

Ejecute el código siguiente en un cuaderno o editor de SQL para crear la función:

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 esta herramienta de recuperación en el agente de IA, se debe encapsular con UCFunctionToolkit. Esto permite el seguimiento automático a través de MLflow.

El rastreo de MLflow captura información detallada de ejecución para aplicaciones de inteligencia artificial generativa. Registra entradas, salidas y metadatos para cada paso, lo que le ayuda a depurar problemas y analizar el rendimiento.

Al usar UCFunctionToolkit, los recuperadores de datos generan automáticamente tipos de extensión RETRIEVER en los registros de MLflow si su salida se ajusta al esquema del recuperador de MLflow. Consulte el Esquema de trazado de MLflow.

Para obtener más información sobre UCFunctionToolkit, consulte la Documentación de Unity Catalog.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Esta herramienta de recuperación tiene las siguientes advertencias:

  • Los clientes SQL pueden limitar el número máximo de filas o bytes devueltos. Para evitar el truncamiento de datos, debe truncar los valores de columna devueltos por la UDF. Por ejemplo, podría usar substring(chunked_text, 0, 8192) para reducir el tamaño de las columnas de contenido de gran tamaño y evitar el truncamiento de filas durante la ejecución.
  • Dado que esta herramienta es un contenedor para la función vector_search(), está sujeto a las mismas limitaciones que la función vector_search(). Consulte Limitaciones.

Si este ejemplo no es adecuado para su caso de uso, cree una herramienta de recuperación de búsqueda vectorial mediante código de agente personalizado en su lugar.

Recuperador de búsqueda vectorial con código de agente (PyFunc)

En el ejemplo siguiente se crea un recuperador de vector de búsqueda para un agente basado en PyFunc dentro del código del agente.

En este ejemplo se usa databricks-vectorsearch para crear un recuperador básico que realice una búsqueda de similitud de vectores de con filtros. Usa decoradores de MLflow para habilitar el seguimiento de agentes.

Nota:

Para ajustarse al esquema del recuperador de MLflow, la función retriever debe devolver un tipo Document y usar el campo metadata de la clase Document para agregar atributos adicionales al documento devuelto, como like doc_uri y similarity_score.

Use el código siguiente en el módulo del agente o en el cuaderno del 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 ejecutar el recuperador, ejecute el siguiente código de Python. Opcionalmente, puede incluir filtros de búsqueda vectorial en la solicitud para filtrar los resultados.

retriever = VectorSearchRetriever()
query = "What is Databricks?"
filters={"text_column LIKE": "Databricks"},
results = retriever(query, filters=filters, score_threshold=0.1)

Configurar esquema del recuperador

Para asegurarse de que los recuperadores se rastreen adecuadamente y se representen correctamente en las aplicaciones de bajada, haga una llamada a mlflow.models.set_retriever_schema al definir su agente. Use set_retriever_schema para asignar los nombres de columna de la tabla devuelta a los campos esperados de MLflow, como primary_key, text_columny 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"],
)

También puede especificar columnas adicionales en el esquema del recuperador proporcionando una lista de nombres de columna con el campo other_columns.

Si tiene varios recuperadores, puede definir varios esquemas mediante nombres únicos para cada esquema del recuperador.

El esquema del recuperador establecido durante la creación del agente afecta a las aplicaciones y flujos de trabajo posteriores, como la aplicación de revisión y los conjuntos de evaluación. En concreto, la columna doc_uri actúa como identificador principal para los documentos devueltos por el recuperador.

  • La aplicación de revisión ,, muestra el doc_uri para ayudar a los revisores a evaluar las respuestas y rastrear los orígenes de los documentos. Consulte Revisión de la interfaz de usuario de la aplicación.
  • Los Conjuntos de evaluación utilizan doc_uri para comparar los resultados del recuperador con conjuntos de datos de evaluación predefinidos para determinar la exhaustividad y precisión del recuperador. Consulte Conjuntos de evaluación.

Seguimiento del recuperador

El seguimiento de MLflow agrega observabilidad mediante la captura de información detallada sobre la ejecución del agente. Proporciona una manera de registrar las entradas, salidas y metadatos asociados a cada paso intermedio de una solicitud, lo que le permite identificar rápidamente el origen de errores y comportamientos inesperados.

En este ejemplo, se utiliza el decorador @mlflow.trace para crear un seguimiento para el recuperador y el analizador. Para ver otras opciones para configurar métodos de seguimiento, consulte Seguimiento de MLflow para agentes.

El decorador crea un intervalo de que se inicia cuando se invoca la función y finaliza cuando la función retorna. MLflow registra automáticamente la entrada y salida de la función y las excepciones generadas.

Nota:

Los usuarios de LangChain, LlamaIndex y la biblioteca de OpenAI pueden usar el registro automático de MLflow en lugar de definir manualmente trazas con el decorador. Consulte Uso del registro automático para agregar seguimientos a los agentes.

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

Para asegurarse de que las aplicaciones de bajada, como Agent Evaluation y AI Playground, representen el seguimiento del recuperador correctamente, asegúrese de que el decorador cumpla los siguientes requisitos:

  • Use span_type="RETRIEVER" y asegúrese de que la función devuelve List[Document] objeto. Consulte Intervalos del recuperador.
  • El nombre de seguimiento y el nombre del retriever_schema deben coincidir para configurar el seguimiento correctamente.

Filtrar resultados de la búsqueda vectorial

Puede limitar el ámbito de búsqueda a un subconjunto de datos mediante un filtro de búsqueda vectorial.

El parámetro filters en VectorSearchRetriever define las condiciones del filtro usando la especificación del filtro de vector de búsqueda de Databricks.

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

Dentro del método __call__, el diccionario de filtros se pasa directamente a la función similarity_search:

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

Después del filtrado inicial, el parámetro score_threshold proporciona filtrado adicional estableciendo una puntuación de similitud mínima.

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

El resultado final incluye documentos que cumplen las condiciones de filters y score_threshold.

Pasos siguientes

Después de crear una herramienta de agente de funciones para el Catálogo de Unity, agregue la herramienta a un agente de IA. Consulte Incorporación de herramientas de Unity Catalog a agentes.