Iscrivere e registrare gli agenti di intelligenza artificiale

Articolo
02/18/2025

Importante

Questa funzionalità è disponibile in anteprima pubblica.

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.

Databricks consiglia di installare la versione più recente del databricks-sdk.

% pip install databricks-sdk

Registrazione basata su codice

Databricks consiglia di usare la funzionalità "Models from Code" di MLflow quando si registrano gli agenti.

In questo approccio il codice dell'agente viene acquisito come file Python e 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.

È possibile associare questo approccio con l'uso di API di convalida pre-distribuzione come mlflow.models.predict() per garantire che l'agente venga eseguito in modo affidabile quando viene distribuito per la gestione.

Il codice che registra l'agente o l'agente deve trovarsi in un notebook separato dal 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

Durante la registrazione, è necessario definire una Signature del Modello MLflow (), che specifica lo schema di input e output dell'agente (). La firma convalida gli input e gli output per garantire che l'agente interagisca correttamente con strumenti downstream come AI Playground e l'app di revisione. Guida anche altre applicazioni su come usare l'agente in modo efficace.

Databricks consiglia di usare le funzionalità di inferenza della firma del modello di MLflow per generare automaticamente la firma dell'agente in base a un esempio di input fornito. 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 seguenti illustrano come registrare un agente con LangChain.

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.
Includi mlflow.models.set_model(lc_agent) nel notebook o file.
Creare un nuovo Notebook da usare come Notebook del driver (in questo esempio, chiamato driver.py).
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.
Distribuire il modello. Vedere Distribuire un agente per un'applicazione di IA generativa.
Quando viene caricato l'ambiente di servizio, agent.py viene eseguito.
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 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)

registrazione basata su codice con PyFunc

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

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.
Includere mlflow.models.set_model(PyFuncClass) nel Notebook o nel file.
Creare un nuovo Notebook da usare come Notebook del driver (in questo esempio, chiamato driver.py).
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 altre informazioni, vedere Specificare le risorse per il pass-through di autenticazione automatica.
Distribuire il modello. Vedere Distribuire un agente per un'applicazione di IA generativa.
Quando viene caricato l'ambiente di servizio, agent.py viene eseguito.
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 più comuni di Databricks, Databricks supporta e consiglia di dichiarare le dipendenze delle risorse per l'agente in anticipo, durante la registrazione. Ciò consente pass-through di autenticazione automatica quando si distribuisce l'agente: Databricks esegue automaticamente il provisioning, ruota e gestisce le credenziali di breve durata per accedere in modo 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.

Nota

Se non si specificano risorse durante la registrazione degli agenti LangChain usando mlflow.langchain.log_model(...), MLflow esegue l'inferenza automatica delle risorse con il massimo sforzo. Tuttavia, ciò potrebbe non includere tutte le dipendenze, causando errori di autorizzazione quando si serve o si esegue una query sull'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.

Tipo di risorsa	Versione minima `mlflow` necessaria per registrare la risorsa
Indice di ricerca vettoriale	Richiede la versione `mlflow` 2.13.1 o superiore
Endpoint di gestione del modello	Richiede la versione `mlflow` 2.13.1 o superiore
Magazzino SQL	Richiede `mlflow` 2.16.1 o versione successiva
Funzione del catalogo Unity	Richiede `mlflow` 2.16.1 o versione successiva
Spazio Genie	Richiede `mlflow` 2.17.1 o versione successiva
Tabella del catalogo Unity	Richiede `mlflow` 2.18.0 o versione successiva

Autenticazione automatica per i client OpenAI

Se l'agente usa il client OpenAI, usare Databricks SDK per l'autenticazione automatica durante la distribuzione. Databricks SDK fornisce un wrapper per costruire il client OpenAI con l'autorizzazione configurata automaticamente. Eseguire quanto segue nel notebook:

% pip install databricks-sdk[openai]

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

Specificare quindi l'endpoint Model Serving come parte di resources per l'autenticazione automatica in fase di distribuzione.

Registrare l'agente a Unity Catalog

Prima di distribuire l'agente, è necessario registrare l'agente in Unity Catalog. Registrazione del pacchetto dell'agente come modello nel Catalogo Unity. Di conseguenza, è possibile utilizzare le autorizzazioni di Unity Catalog per gestire le 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)

Condividi tramite