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.

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 tramite codice

Databricks consiglia di usare la funzionalità Models from Code di MLflow durante la registrazione degli 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.

Per un esempio di registrazione basata sul codice, vedere notebook di esempio per la creazione di ChatAgent.

Dedurre la firma del modello durante la registrazione dei log

Nota

Databricks consiglia di creare un agente usando l'interfaccia ChatAgent. Se si usa ChatAgent, è possibile ignorare questa sezione; MLflow deduce automaticamente una firma valida per l'agente.

Se non si usa l'interfaccia ChatAgent, è necessario usare uno dei metodi seguenti per specificare il MLflow Model Signature dell'agente al momento della registrazione:

1. Definire manualmente la firma
2. 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.

La firma del modello MLflow convalida input e 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.

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.

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 Autenticazione per le risorse di Databricks.

5. Distribuire il modello. Vedere Distribuire un agente per applicazioni di intelligenza artificiale generative.

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 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.

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 usare log_model() per 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 Autenticazione per le risorse di Databricks.

5. Distribuire il modello. Vedere Distribuire un agente per applicazioni di intelligenza artificiale generative.

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

Autenticazione per le risorse di Databricks

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.

Model Serving supporta due tipi diversi di autenticazione per le risorse gestite da Databricks:

1. Autenticazione di sistema: consente all'entità servizio agente di accedere a qualsiasi risorsa dipendente specificata in fase di registrazione dell'agente. Ciò è utile per accedere a risorse condivise o non sensibili, ad esempio un indice di ricerca vettoriale contenente la documentazione pubblica
2. [Beta] Autenticazione on-behalf-of-user: consente all'agente di usare le credenziali dell'utente finale per accedere alle risorse di Databricks. Ciò è utile per gli scenari in cui l'agente deve accedere a dati sensibili o eseguire query sulle API remote per eseguire azioni su base utente

Specificare le risorse per il pass-through di autenticazione automatica (autenticazione di sistema)

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 autenticazione passante automatica quando si distribuisce l'agente: Databricks fornisce, ruota e gestisce automaticamente le credenziali di breve durata per garantire accessi sicuri a queste dipendenze delle risorse dal punto di accesso dell'agente.

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

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

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-meta-llama-3-3-70b-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"),
      DatabricksApp(app_name="app_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 interroga l'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
Connessione al catalogo Unity	Richiede mlflow 2.17.1 o versione successiva

Autenticazione per conto dell'utente

Importante

Questa funzionalità è in versione beta.

Quando si registra un agente che utilizza l'autenticazione per conto dell'utente, specificare il set minimo di ambiti API di Databricks necessari per eseguire azioni come l'utente finale all'interno del codice dell'agente. Ciò garantisce che l'agente disponga dell'accesso con privilegi minimi per eseguire azioni per conto dell'utente finale durante la distribuzione, migliorando la sicurezza impedendo azioni non autorizzate e riducendo al minimo il rischio di uso improprio dei token.

Di seguito è riportato un elenco di ambiti necessari per accedere a diversi tipi comuni di risorse di Databricks.

Risorsa Databricks	Ambiti API obbligatori
Indice di ricerca vettoriale	serving.serving-endpoints,vectorsearch.vector-search-endpointsvectorsearch.vector-search-indexes
Endpoint di gestione del modello	serving.serving-endpoints
SQL Warehouse	sql.statement-execution,sql.warehouses
Connessioni UC	catalog.connections
Spazio Genie	dashboards.genie

Per abilitare l'autenticazione on-behalf-of-user, passare un MLflow AuthPolicy a log_model(), come illustrato nell'esempio seguente. Un flusso ML include AuthPolicy due componenti:

• system_auth_policy: specificare le risorse per l'autenticazione di sistema. In genere, gli agenti useranno l'autenticazione di sistema per le risorse condivise (ad esempio per eseguire query sugli endpoint di gestione del modello) in combinazione con l'autenticazione on-behalf-of-user per l'accesso a risorse o API sensibili
• user_auth_policy: Specificare gli ambiti API necessari all'agente per l'autenticazione per conto dell'utente

from mlflow.models.resources import DatabricksServingEndpoint
from mlflow.models.auth_policy import SystemAuthPolicy, UserAuthPolicy, AuthPolicy

resources = [
    DatabricksServingEndpoint(endpoint_name="databricks-meta-llama-3-3-70b-instruct")
]
# Specify resources here for system authentication
system_auth_policy = SystemAuthPolicy(resources=resources)

# Specify the minimal set of API scopes needed for on-behalf-of-user authentication
	# When deployed, the agent can access Databricks resources and APIs
	# on behalf of the end user, but only via REST APIs that are covered by the list of
	# scopes below

user_auth_policy = UserAuthPolicy(
    api_scopes=[
        "serving.serving-endpoints",
        "vectorsearch.vector-search-endpoints",
        "vectorsearch.vector-search-indexes",
    ]
)

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        ...
        # Instead of passing `resources` (which only supports system authentication),
	        # pass an auth_policy to log_model to enable both system authentication and
	        # on-behalf-of-user authentication
        auth_policy=AuthPolicy(
            system_auth_policy=system_auth_policy,
            user_auth_policy=user_auth_policy
        )
    )

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, get_open_ai_client(). 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 su Unity Catalog

Prima di distribuire l'agente, è necessario registrare l'agente in Unity Catalog. Il pacchetto dell'agente viene registrato 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)

Vedi mlflow.register_model().

Passaggi successivi

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