Iscrivere e registrare gli agenti di intelligenza artificiale

Registrare gli agenti di intelligenza artificiale con Mosaic AI Agent Framework. La registrazione di un agente è la base del processo di sviluppo. La registrazione acquisisce un "punto nel tempo" del codice e della configurazione dell'agente, in modo da poter valutare la qualità della configurazione stessa.

Requisiti

Creare un agente di intelligenza artificiale prima di registrarlo.

Confronto tra registrazione basata su codice e basata sulla serializzazione

È possibile usare la registrazione MLflow basata su codice o la registrazione MLflow basata su serializzazione. Databricks consiglia di usare la registrazione basata su codice.

La registrazione MLflow basata su codice: il codice dell'agente viene salvato come file Python. L'ambiente Python viene acquisito come elenco di pacchetti. Quando l'agente viene distribuito, viene ripristinato l'ambiente Python e il codice dell'agente viene eseguito per caricare l'agente in memoria in modo che possa essere richiamato quando viene chiamato l'endpoint.

di registrazione MLflow basata sulla serializzazione: il codice dell'agente e lo stato corrente nell'ambiente Python vengono serializzati su disco, spesso usando librerie come pickle o joblib. Quando l'agente viene distribuito, viene ripristinato l'ambiente Python e l'oggetto serializzato viene caricato in memoria in modo che possa essere richiamato quando viene chiamato l'endpoint.

La tabella mostra i vantaggi e gli svantaggi di ciascun metodo.

Per la registrazione basata su codice, il codice che effettua la registrazione di uno o più agenti deve trovarsi in un notebook separato rispetto al codice dell'agente. Questo Notebook è chiamato Notebook driver. Per un Notebook di esempio, vedere Notebook di esempio.

Dedurre la firma del modello durante la registrazione

Per registrare e distribuire un agente, è necessario registrarlo con unMLflow Model Signature . La firma definisce lo schema di input e output dell'agente.

Databricks consiglia di usare le funzionalità di inferenza della firma del modello di MLflow per determinare automaticamente la firma dell'agente fornendo un esempio di input. Questo approccio è più pratico rispetto alla definizione manuale della firma.

Gli esempi LangChain e PyFunc di seguito usano l'inferenza della firma del modello.

Se si preferisce definire in modo esplicito una firma del modello in fase di registrazione, vedere documentazione di MLflow - Come registrare i modelli con firme.

Registrazione basata su codice con LangChain

Le istruzioni e l'esempio di codice in questa sezione illustrano come registrare un agente con LangChain.

1. Creare un Notebook o un file Python con il codice. Per questo esempio, il notebook o il file è denominato agent.py. Il notebook o il file deve contenere un agente LangChain, indicato qui come lc_agent.

2. Includi mlflow.models.set_model(lc_agent) nel notebook o file.

3. Creare un nuovo Notebook da usare come Notebook del driver (in questo esempio, chiamato driver.py).

4. Nel notebook del driver usare il codice seguente per eseguire agent.py e registrare i risultati in un modello MLflow:

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

   Il parametro resources dichiara le risorse gestite da Databricks necessarie per gestire l'agente, ad esempio un indice di ricerca vettoriale o un endpoint che serve un modello di base. Per altre informazioni, vedere Specificare le risorse per il pass-through di autenticazione automatica.

5. Distribuire il modello. Vedere Distribuire un agente per un'applicazione di IA generativa.

6. Quando viene caricato l'ambiente di servizio, agent.py viene eseguito.

7. Quando arriva una richiesta di servizio, lc_agent.invoke(...) viene richiamato.

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 your 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)

registrazione basata su codice con PyFunc

Le istruzioni e l'esempio di codice in questa sezione illustrano come registrare un agente con PyFunc.

1. Creare un Notebook o un file Python con il codice. Per questo esempio, il notebook o il file è denominato agent.py. Il notebook o il file deve contenere una classe PyFunc, denominata PyFuncClass.

2. Includere mlflow.models.set_model(PyFuncClass) nel Notebook o nel file.

3. Creare un nuovo Notebook da usare come Notebook del driver (in questo esempio, chiamato driver.py).

4. Nel notebook del driver usare il codice seguente per eseguire agent.py e registrare i risultati in un modello MLflow:

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

   Il parametro resources dichiara le risorse gestite da Databricks necessarie per gestire l'agente, ad esempio un indice di ricerca vettoriale o un endpoint che serve un modello di base. Per ulteriori informazioni, vedere Specificare le risorse per il passaggio di autenticazione automatica.

5. Distribuire il modello. Vedere Distribuire un agente per un'applicazione di IA generativa.

6. Quando viene caricato l'ambiente di servizio, agent.py viene eseguito.

7. Quando arriva una richiesta di servizio, PyFuncClass.predict(...) viene richiamato.

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)

Specificare le risorse per il pass-through di autenticazione automatica

Gli agenti di intelligenza artificiale spesso devono eseguire l'autenticazione ad altre risorse per completare le attività. Ad esempio, un agente potrebbe dover accedere a un indice di ricerca vettoriale per eseguire query su dati non strutturati.

Come descritto in Authentication for dependent resources, Model Serving supporta l'autenticazione sia per le risorse gestite da Databricks che per le risorse esterne quando si distribuisce l'agente.

Per i tipi di risorse Databricks più comuni, Databricks supporta e consiglia di dichiarare in anticipo le dipendenze delle risorse dell'agente durante il logging. Ciò consente l'autenticazione automatica di passaggio quando si distribuisce l'agente: Databricks fornisce automaticamente, ruota e gestisce credenziali a breve termine per consentire l'accesso sicuro a queste dipendenze delle risorse dall'endpoint dell'agente.

Per abilitare il pass-through di autenticazione automatica, specificare le risorse dipendenti usando il parametro resources dell'API log_model(), come illustrato nel codice seguente.

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

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"),
    ]
  )

Databricks consiglia di specificare manualmente resources per tutte le versioni dell'agente.

La tabella seguente elenca le risorse di Databricks che supportano il pass-through di autenticazione automatica e la versione minima mlflow necessaria per registrare la risorsa.

Registrare l'agente a Unity Catalog

Prima di distribuire l'agente, è necessario registrare l'agente in Unity Catalog. Quando si registra l'agente, viene impacchettato come modello nel Catalogo Unity ed è possibile usare le autorizzazioni del Catalogo Unity per l'autorizzazione delle risorse nell'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)

Passaggi successivi

• Aggiungere tracce a un agente di intelligenza artificiale.
• Distribuzione di un agente di IA.