Creación y rastreo de herramientas de recuperación para datos no estructurados

Use mosaic AI Agent Framework para crear herramientas que permitan a los agentes de IA consultar datos no estructurados, como una colección de documentos. En esta página se muestra cómo:

• Desarrollar recuperadores localmente
• Creación de un recuperador mediante funciones de Catálogo de Unity
• Consulta de índices vectoriales externos
• Agregar rastreo MLflow para observabilidad

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

Desarrollo local de herramientas de búsqueda de vectores con AI Bridge

La forma más rápida de empezar a crear una herramienta de recuperación de Búsqueda vectorial de Databricks es desarrollarla y probarla localmente mediante paquetes de Databricks AI Bridge como databricks-langchain y databricks-openai.

LangChain/LangGraph

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

%pip install --upgrade databricks-langchain

El código siguiente prototipa una herramienta de recuperación que consulta un índice hipotético de búsqueda vectorial y lo enlaza localmente a un LLM para que pueda probar su comportamiento de llamada a herramientas.

Proporcione un descriptivo tool_description para ayudar al agente a comprender la herramienta y determinar cuándo invocarla.

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-sonnet-4-5")
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 escenarios que usan índices de acceso directo o índices Delta Sync utilizando incrustaciones autogestionadas, debe configurar VectorSearchRetrieverTool y especificar un modelo de incrustación personalizado y una columna de texto. Consulte las opciones para proporcionar incrustaciones.

En el ejemplo siguiente se muestra cómo configurar un VectorSearchRetrieverTool con claves columns y 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 más información, consulte los documentos de API para VectorSearchRetrieverTool.

OpenAI

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

%pip install --upgrade databricks-openai

El código siguiente crea prototipos de un recuperador que consulta un índice hipotético de búsqueda vectorial y lo integra con los modelos GPT de OpenAI.

Proporcione un descriptivo tool_description para ayudar al agente a comprender la herramienta y determinar cuándo invocarla.

Para más información sobre recomendaciones OpenAI para herramientas, ver Documentación OpenAI Function Calling.

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 escenarios que usan índices de acceso directo o índices Delta Sync utilizando incrustaciones autogestionadas, debe configurar VectorSearchRetrieverTool y especificar un modelo de incrustación personalizado y una columna de texto. Consulte las opciones para proporcionar incrustaciones.

En el ejemplo siguiente se muestra cómo configurar un VectorSearchRetrieverTool con claves columns y 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 más información, consulte los documentos de API para VectorSearchRetrieverTool.

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.

En la sección siguiente se muestra cómo migrar el recuperador a una función de catálogo de Unity.

Herramienta de recuperación de búsqueda vectorial con funciones de Unity Catalog

Puede crear una función de Unity Catalog que envuelva una consulta de índice Mosaic AI Vector Search. Este enfoque:

• Admite casos de uso de producción con gobernanza y capacidad de descubrimiento
• Usa la función SQL vector_search() en segundo plano
• Admite el seguimiento automático de MLflow.
  • Debe alinear la salida de la función con el esquema del recuperador de MLflow mediante el uso de los alias page_content y metadata.
  • Cualquier columna de metadatos adicional debe añadirse a la columna metadata utilizando la función SQL map, 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 rastreo automático a través de MLflow generando automáticamente RETRIEVER tipos de span en los registros de MLflow.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

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

Las herramientas de recuperación del catálogo de Unity tienen las siguientes limitaciones:

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

Para obtener más información sobre UCFunctionToolkit consulte Documentación del catálogo de unidad.

Recuperador que consulta un índice vectorial hospedado fuera de Databricks

Si el índice vectorial se hospeda fuera de Azure Databricks, puede crear una conexión de Unity Catalog para conectarse al servicio externo y usar la conexión en el código de su agente. Consulte Conexión de herramientas de agente de IA a servicios externos.

El siguiente ejemplo crea un recuperador que llama a un índice vectorial hospedado fuera de Databricks para un agente con sabor a PyFunc.

1. Cree una conexión de catálogo de Unity al servicio externo, en este caso, 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 la herramienta retriever en el código del agente usando la conexión Unity Catalog. En este ejemplo se usan decoradores de MLflow para habilitar el seguimiento del agente.

   Nota:

   Para ajustarse al esquema del recuperador de MLflow, la función retriever debe devolver un List[Document] objeto y usar el metadata campo de la clase Document para agregar atributos adicionales al documento devuelto, como doc_uri y similarity_score. Consulte el documento de 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 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 = [0.01944167, 0.0040178085 . . .  TRIMMED FOR BREVITY 010858015, -0.017496133]
   results = retriever(query, score_threshold=0.1)
   

Añada seguimiento a un recuperador

Añada seguimiento MLflow para monitorizar y depurar su recuperador. El seguimiento le permite ver entradas, salidas y metadatos para cada paso de ejecución.

El ejemplo anterior añade el decorador @mlflow.trace a los métodos __call__ y análisis. 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.

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 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 (https://mlflow.org/docs/latest/tracing/tracing-schema.html#retriever-spans) y asegúrese de que la función devuelve un objeto List[Document].
• El nombre de seguimiento y el retriever_schema nombre deben coincidir para configurar el seguimiento correctamente. Consulte la siguiente sección para conocer cómo configurar el esquema del sistema de recuperación.

Establecer el esquema del recuperador para asegurar la compatibilidad con MLflow

Si el seguimiento devuelto desde el recuperador o span_type="RETRIEVER" no se ajusta al esquema del recuperador estándar de MLflow, debe asignar manualmente el esquema devuelto a MLflow campos esperados. Esto garantiza que MLflow pueda rastrear correctamente su recuperador y representar trazas en aplicaciones posteriores.

Para establecer el esquema del recolector manualmente:

1. Llame a mlflow.models.set_retriever_schema cuando defina 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"],
   )
   

2. Especifique columnas adicionales en el esquema del recuperador proporcionando una lista de nombres de columna con el other_columns campo .

3. 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 (MLflow 2).

Pasos siguientes

Después de crear una herramienta de catálogo de Unity, agréguela al agente. Consulte Creación de una herramienta de agente.