Comparteix via

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 "momento dado" del código y la configuración del agente para que pueda evaluar la calidad de la configuración.

Requirements

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 de Modelos desde código 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 para que se pueda 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 Los cuadernos de ejemplo de creación de ResponsesAgent.

Deducción de la firma del modelo durante el registro

Note

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

Si no usa la interfaz de ResponsesAgent, 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.

Los ejemplos de LangChain y PyFunc siguientes usan la inferencia de la 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 Implementar el paso automático de autenticación.

  5. Se implementa el modelo. Consulte Implementación de un agente para aplicaciones de IA generativas.

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

  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 usar log_model() para 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 Implementar el paso automático de autenticación.

  5. Se implementa el modelo. Consulte Implementación de un agente para aplicaciones de IA generativas.

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

  7. Cuando se incluye una solicitud de servicio, se llama 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-meta-llama-3-3-70b-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)

Autenticación para recursos de Databricks

A menudo, los agentes de inteligencia artificial necesitan autenticarse en otros recursos para completar tareas. Por ejemplo, un agente implementado podría necesitar acceder a un índice de búsqueda vectorial para consultar datos no estructurados o acceder al Registro de mensajes para cargar mensajes dinámicos.

El paso de autenticación automática y la autenticación en nombre de requieren configuración durante el registro del agente.

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)

Consulte mlflow.register_model().

Pasos siguientes

  • Last updated on