Registro y registro de agentes de IA

Registro de agentes de IA mediante Mosaic AI Agent Framework. Registrar un agente es la base del proceso de desarrollo. El registro captura un "punto en el tiempo" del código y la configuración del agente para permitir evaluar la calidad de la configuración.

Requisitos

Cree un agente de IA antes de registrarlo.

Databricks recomienda instalar la versión más reciente del databricks-sdk.

% pip install databricks-sdk

Registro basado en código

Databricks recomienda usar la funcionalidad "Models from Code" de MLflow al registrar agentes.

En este enfoque, el código del agente se captura como un archivo de Python y el entorno de Python se captura como una lista de paquetes. Cuando se implementa el agente, se restaura el entorno de Python y el código del agente se ejecuta para cargar el agente en la memoria, por lo que se puede invocar cuando se llama al punto de conexión.

Puede acoplar este enfoque con el uso de API de validación previa a la implementación, como mlflow.models.predict() para asegurarse de que el agente se ejecuta de forma confiable cuando se implementa para servir.

Para ver un ejemplo de registro basado en código, consulte cuadernos de ejemplo de creación de ChatAgent.

Deducción de la firma del modelo durante el registro

Nota:

Databricks recomienda crear un agente mediante la interfaz ChatAgent. Si usa ChatAgent, puede omitir esta sección; MLflow deduce automáticamente una firma válida para el agente.

Si no usa la interfaz de ChatAgent, debe usar uno de los métodos siguientes para especificar el MLflow de MLflow Firma de modelo en tiempo de registro:

1. Definir manualmente la firma
2. Use las funcionalidades de inferencia de firma de modelos de MLflow para generar automáticamente la firma del agente en función de un ejemplo de entrada que proporcione. Este enfoque es más conveniente que definir manualmente la firma.

La firma del modelo de MLflow valida las entradas y salidas para asegurarse de que el agente interactúa correctamente con herramientas de nivel inferior, como AI Playground y la aplicación de revisión. También guía otras aplicaciones sobre cómo usar el agente de forma eficaz.

En los ejemplosLangChain y PyFunc siguientes se usa la inferencia de firma del modelo.

Si prefiere definir explícitamente una firma de modelo en el momento del registro, consulte los documentos de MLflow: Cómo registrar modelos con firmas.

Registro basado en código con LangChain

Las instrucciones y el ejemplo de código siguientes muestran cómo registrar un agente con LangChain.

1. Cree un cuaderno o un archivo de Python con el código. En este ejemplo, el cuaderno o el archivo se denominan agent.py. El cuaderno o el archivo deben contener un agente langChain, que se conoce aquí como lc_agent.

2. Incluya mlflow.models.set_model (lc_agent) en el cuaderno o archivo.

3. Cree un cuaderno que actúe como cuaderno de controlador (denominado driver.py en este ejemplo).

4. En el cuaderno del controlador, use el código siguiente para ejecutar agent.py y registrar los resultados en un modelo de MLflow:

   mlflow.langchain.log_model(lc_model="/path/to/agent.py", resources=list_of_databricks_resources)
   

   El parámetro resources declara los recursos administrados por Databricks necesarios para atender al agente, como un índice de búsqueda vectorial o un punto de conexión de servicio que sirve un modelo de base. Para obtener más información, consulte Especificar recursos para la autenticación automática.

5. Se implementa el modelo. Consulte Implementación de un agente para una aplicación de IA generativa.

6. Cuando se carga el entorno de servicio, se ejecuta agent.py.

7. Cuando se incluye una solicitud de servicio, se llama a lc_agent.invoke(...).

import mlflow

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
  "messages": [
    {
      "role": "user",
      "content": "What is Retrieval-augmented Generation?",
    }
  ]
}

# example using langchain
with mlflow.start_run():
  logged_agent_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this configuration is used by agent code. The development_config is overwritten.
    artifact_path="agent", # This string is used as the path inside the MLflow model where artifacts are stored
    input_example=input_example, # Must be a valid input to the agent
    example_no_conversion=True, # Required
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.langchain.load_model(logged_agent_info.model_uri)
model.invoke(example)

Registro basado en código con PyFunc

Las instrucciones y el ejemplo de código siguientes muestran cómo registrar un agente con PyFunc.

1. Cree un cuaderno o un archivo de Python con el código. En este ejemplo, el cuaderno o el archivo se denominan agent.py. El cuaderno o el archivo deben contener una clase PyFunc, denominada PyFuncClass.

2. Incluya mlflow.models.set_model(PyFuncClass) en el cuaderno o el archivo.

3. Cree un cuaderno que actúe como cuaderno de controlador (denominado driver.py en este ejemplo).

4. En el cuaderno del controlador, use el código siguiente para ejecutar agent.py y registrar los resultados en un modelo de MLflow:

   mlflow.pyfunc.log_model(python_model="/path/to/agent.py", resources=list_of_databricks_resources)
   

   El parámetro resources declara los recursos administrados por Databricks necesarios para atender al agente, como un índice de búsqueda vectorial o un punto de conexión de servicio que sirve un modelo de base. Para obtener más información, consulte Especificar recursos para la autenticación automática.

5. Se implementa el modelo. Consulte Implementación de un agente para una aplicación de IA generativa.

6. Cuando se carga el entorno de servicio, se ejecuta agent.py.

7. Cuando llega una petición de servicio, se invoca a PyFuncClass.predict(...).

import mlflow
from mlflow.models.resources import (
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex,
)

code_path = "/Workspace/Users/first.last/agent.py"
config_path = "/Workspace/Users/first.last/config.yml"

# Input example used by MLflow to infer Model Signature
input_example = {
  "messages": [
    {
      "role": "user",
      "content": "What is Retrieval-augmented Generation?",
    }
  ]
}

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    artifact_path="agent",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
    ]
  )

print(f"MLflow Run: {logged_agent_info.run_id}")
print(f"Model URI: {logged_agent_info.model_uri}")

# To verify that the model has been logged correctly, load the agent and call `invoke`:
model = mlflow.pyfunc.load_model(logged_agent_info.model_uri)
model.invoke(example)

Especificar recursos para la transferencia automática de autenticación

A menudo, los agentes de inteligencia artificial deben autenticarse en otros recursos para completar tareas. Por ejemplo, un agente puede necesitar acceder a un índice de búsqueda vectorial para consultar datos no estructurados.

Como se describe en Autenticación para recursos dependientes, Model Serving admite la autenticación en recursos administrados por Databricks y externos al implementar el agente.

Para los tipos de recursos de Databricks más comunes, Databricks admite y recomienda declarar las dependencias de recursos para el agente por adelantado durante el registro. Esto permite paso de autenticación automática al implementar el agente: Databricks aprovisiona, rota y administra automáticamente las credenciales de corta duración para acceder de forma segura a estas dependencias de recursos desde el punto de conexión del agente.

Para habilitar el paso directo de autenticación automática, especifique los recursos dependientes mediante el parámetro resources de la API log_model(), como se muestra en el código siguiente.

import mlflow
from mlflow.models.resources import (
  DatabricksVectorSearchIndex,
  DatabricksServingEndpoint,
  DatabricksSQLWarehouse,
  DatabricksFunction,
  DatabricksGenieSpace,
  DatabricksTable,
  DatabricksUCConnection
)

with mlflow.start_run():
  logged_agent_info = mlflow.pyfunc.log_model(
    python_model=agent_notebook_path,
    artifact_path="agent",
    input_example=input_example,
    example_no_conversion=True,
    # Specify resources for automatic authentication passthrough
    resources=[
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
      DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
      DatabricksFunction(function_name="ml.tools.python_exec"),
      DatabricksGenieSpace(genie_space_id="your_genie_space_id"),
      DatabricksTable(table_name="your_table_name"),
      DatabricksUCConnection(connection_name="your_connection_name"),
    ]
  )

Databricks recomienda especificar manualmente resources para todos los tipos de agente.

Nota:

Si no especifica recursos al registrar agentes de LangChain mediante mlflow.langchain.log_model(...), MLflow realiza la inferencia automática de los recursos en la medida de lo posible. Sin embargo, esto puede no capturar todas las dependencias, lo que da lugar a errores de autorización al atender o consultar al agente.

En la tabla siguiente se enumeran los recursos de Databricks que admiten el paso a través de la autenticación automática y la versión mínima mlflow necesaria para registrar el recurso.

Tipo de recurso	Versión mínima mlflow necesaria para registrar el recurso
Índice de búsqueda vectorial	Requiere mlflow 2.13.1 o superior
Modelo de puntos de conexión de servicio	Requiere mlflow 2.13.1 o superior
Almacén de SQL	Requiere mlflow 2.16.1 o superior
La función de catálogo de Unity	Requiere mlflow 2.16.1 o superior
Espacio de Genie	Requiere mlflow 2.17.1 o superior
Tabla de catálogos de Unity	Requiere mlflow 2.18.0 o superior
Conexión a Unity Catalog	Requiere mlflow 2.17.1 o superior

Autenticación automática para clientes openAI

Si el agente usa el cliente openAI, use el SDK de Databricks para autenticarse automáticamente durante la implementación. El SDK de Databricks proporciona un contenedor para construir el cliente openAI con autorización configurada automáticamente. Ejecute lo siguiente en el cuaderno:

% pip install databricks-sdk[openai]

from databricks.sdk import WorkspaceClient
def openai_client(self):
  w = WorkspaceClient()
  return w.serving_endpoints.get_open_ai_client()

A continuación, especifique el punto de conexión De servicio de modelos como parte de para autenticarse automáticamente en el momento de resources la implementación.

Registrar el agente en el catálogo de Unity

Antes de implementar el agente, debe registrarlo en Unity Catalog. Al registrar el agente se empaqueta como un modelo en Unity Catalog. Como resultado, puede usar permisos de Catálogo de Unity para la autorización de los recursos del agente.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "agent_name"

model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=model_name)

Pasos siguientes

• Adición de trazados a un agente de IA.
• Implementar un agente de IA.