Registro de agentes de IA

Artículo
09/27/2024

Importante

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

Registro de agentes de IA mediante Mosaic AI Agent Framework. Registrar un agente es la base del proceso de desarrollo. Al registrar se captura un "momento dado" del código y la configuración del agente para que pueda evaluar la calidad de la configuración.

Requisitos

Cree un agente de IA antes de registrarlo.

Registro basado en código frente al registro basado en serialización

Puede usar el registro de MLflow basado en código o el registro de MLflow basado en serialización. Databricks recomienda usar el registro basado en código.

Registro de MLflow basado en código: el código de la cadena se captura como archivo Python. El entorno de Python se captura como lista de paquetes. Cuando se implementa la cadena, se restaura el entorno de Python y el código de la cadena se ejecuta para cargarla en la memoria, por lo que se puede invocar cuando se llama al punto de conexión.

Registro de MLflow basado en serialización: el código de la cadena y el estado actual del entorno de Python se serializan en el disco, para lo que a menudo se usan bibliotecas como pickle o joblib. Cuando se implementa la cadena, se restaura el entorno de Python y el objeto serializado se carga en la memoria, por lo que se puede invocar cuando se llama al punto de conexión.

La tabla muestra las ventajas y desventajas de cada método.

Método	Ventajas	Inconvenientes
Registro de MLflow basado en código	- Supera las limitaciones inherentes de la serialización, no compatible con muchas bibliotecas populares de GenAI. - Guarda una copia del código original para consultarla más adelante. - No es necesario reestructurar el código en un solo objeto que se pueda serializar.	Debe llamarse a `log_model(...)` desde un cuaderno distinto al del código de la cadena (llamado cuaderno de controlador).
Registro de MLflow basado en serialización	Puede llamarse a `log_model(...)` desde el mismo cuaderno donde se define el modelo.	- El código original no está disponible. - Todas las bibliotecas y objetos usados en la cadena deben admitir la serialización.

Para el registro basado en código, el código que registra el agente o la cadena debe estar en un cuaderno independiente del código de la cadena. Este cuaderno se denomina cuaderno de controlador. Para ver un cuaderno de ejemplo, consulte Cuadernos de ejemplo.

Registro basado en código con LangChain

Cree un cuaderno o un archivo de Python con el código. Para los fines de este ejemplo, el cuaderno o el archivo se denomina chain.py. El cuaderno o el archivo deben contener una cadena LangChain, aquí denominada lc_chain.
Incluya mlflow.models.set_model(lc_chain) en el cuaderno o el archivo.
Cree un cuaderno que actúe como cuaderno de controlador (denominado driver.py en este ejemplo).
En el cuaderno del controlador, use mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”)para ejecutar chain.py y registrar los resultados en un modelo de MLflow.
Se implementa el modelo. Consulte Implementación de un agente para una aplicación de IA generativa. La implementación del agente puede depender de otros recursos de Databricks, como un índice de búsqueda vectorial y un modelo que atiende los puntos de conexión. Para los agentes de LangChain:
- MLflow log_model deduce las dependencias requeridas por la cadena y las registra en el archivo MLmodel en el artefacto del modelo registrado.
- Durante la implementación, databricks.agents.deploy crea automáticamente los tokens de OAuth de M2M necesarios para acceder a estas dependencias de recursos inferidas y comunicarse con ellas.
Cuando se carga el entorno de servicio, se ejecuta chain.py.
Cuando se recibe una solicitud de servicio, se llama a lc_chain.invoke(...).


import mlflow

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

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using LangChain
with mlflow.start_run():
  logged_chain_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this is the configuration that is used for training the model. The development_config is overwritten.
    artifact_path="chain", # 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 your chain
    example_no_conversion=True, # Required
  )

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

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

Registro basado en código con PyFunc

Cree un cuaderno o un archivo de Python con el código. Para los fines de este ejemplo, el cuaderno o el archivo se denomina chain.py. El cuaderno o el archivo deben contener una clase PyFunc, aquí denominada PyFuncClass.
Incluya mlflow.models.set_model(PyFuncClass) en el cuaderno o el archivo.
Cree un cuaderno que actúe como cuaderno de controlador (denominado driver.py en este ejemplo).
En el cuaderno del controlador, use mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”, resources=”/path/to/resources.yaml”) para ejecutar chain.py y registrar los resultados en un modelo de MLflow. El parámetro resources declara los recursos necesarios para servir al modelo como, por ejemplo, un índice de búsqueda vectorial o un punto de conexión de servicio que sirve a un modelo de base. Vea un archivo de recursos de ejemplo para PyFunc.
Se implementa el modelo. Consulte Implementación de un agente para una aplicación de IA generativa.
Cuando se carga el entorno de servicio, se ejecuta chain.py.
Cuando se recibe una solicitud de servicio, se llama a PyFuncClass.predict(...).

import mlflow

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

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using PyFunc model

resources_path = "/Workspace/Users/first.last/resources.yml"

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
  )

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

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

Especificar recursos para el agente pyFunc

Puede especificar recursos, como un índice de búsqueda vectorial y un punto de conexión de servicio, que son necesarios para servir el modelo. Para LangChain, los recursos se seleccionan y registran automáticamente junto con el modelo.

Al implementar un agente de tipo pyfunc, debe agregar manualmente las dependencias de recursos del agente implementado. Se crea un token de OAuth de M2M con acceso a todos los recursos especificados del resources parámetro y se proporciona al agente implementado.

Nota:

Para invalidar los recursos a los que tiene permiso el punto de conexión, especifique manualmente los recursos al registrar la cadena.

A continuación se muestra cómo agregar dependencias de índice de búsqueda vectorial y de punto de conexión de servicio especificándolas en el parámetro resources.

 with mlflow.start_run():
   logged_chain_info = mlflow.pyfunc.log_model(
     python_model=chain_notebook_path,
     artifact_path="chain",
     input_example=input_example,
     example_no_conversion=True,
     resources=[
            DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
            DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
            DatabricksVectorSearchIndex(index_name="rag.studio_bugbash.databricks_docs_index")
        ]
   )

También puede agregar recursos si los especifica en un archivo resources.yaml. Puede hacer referencia a esa ruta de acceso del archivo en el parámetro resources. Se crea un token de OAuth de M2M con acceso a todos los recursos especificados en el resources.yaml y se proporciona al agente implementado.

A continuación se muestra un archivo de ejemplo resources.yaml que define los puntos de conexión de servicio del modelo y un índice de búsqueda vectorial.


api_version: "1"
databricks:
  vector_search_index:
    - name: "catalog.schema.my_vs_index"
  serving_endpoint:
    - name: databricks-dbrx-instruct
    - name: databricks-bge-large-en

Registro de la cadena en Unity Catalog

Antes de implementar la cadena, debe registrarla en Unity Catalog. Al registrar la cadena, se empaqueta como modelo en Unity Catalog y puede usar los permisos de Unity Catalog para la autorización de los recursos de la cadena.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "chain_name"

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

Compartir a través de