Pular para o conteúdo principal

registro e registro AI agentes

AI Registre os agentes Mosaic AI do usando o Agent Framework. Registrar um agente é a base do processo de desenvolvimento. O registro captura um “ponto no tempo” do código e da configuração do agente para que você possa avaliar a qualidade da configuração.

Requisitos

Crie um agente AI antes de registrá-lo.

A Databricks recomenda a instalação da versão mais recente do databricks-sdk.

Python
% pip install databricks-sdk

Registro baseado em código

A Databricks recomenda o uso da funcionalidade Models from Code do MLflow ao registrar agentes.

Nessa abordagem, o código do agente é capturado como um arquivo Python e o ambiente Python é capturado como uma lista de pacote. Quando o agente é implantado, o ambiente Python é restaurado e o código do agente é executado para carregar o agente na memória para que ele possa ser invocado quando o endpoint for chamado.

O senhor pode combinar essa abordagem com o uso da validação de pré-implantação APIs como mlflow.models.predict() para garantir que o agente seja executado de forma confiável quando implantado para servir.

Para ver um exemplo de registro baseado em código, consulte o exemplo de criação do ChatAgent Notebook.

Inferir a assinatura do modelo durante o registro

nota

A Databricks recomenda a criação de um agente usando a interface ChatAgent. Se estiver usando o ChatAgent, o senhor pode pular esta seção; o MLflow infere automaticamente uma assinatura válida para o seu agente.

Se não estiver usando a interface ChatAgent, o senhor deve usar um dos métodos a seguir para especificar a assinatura do modelo MLflow do agente no momento do registro em log:

  1. Defina manualmente a assinatura
  2. Use os recursos de inferência de assinatura de modelo do MLflow para gerar automaticamente a assinatura do agente com base em um exemplo de entrada fornecido pelo senhor. Essa abordagem é mais conveniente do que definir manualmente a assinatura.

A assinatura do modelo do MLflow valida entradas e saídas para garantir que o agente interaja corretamente com as ferramentas de downstream do , como o AI Playground e o aplicativo de avaliação. Ele também orienta outros aplicativos sobre como usar o agente de forma eficaz.

Os exemplos de LangChain e PyFunc abaixo usam a inferência de assinatura de modelo.

Se o senhor preferir definir explicitamente uma assinatura de modelo no momento do registro, consulte MLflow docs - How to modelos logged with signatures.

Registro de log baseado em código com LangChain

As instruções e o exemplo de código a seguir mostram ao senhor como log um agente com LangChain.

  1. Crie um arquivo do Notebook ou Python com seu código. Neste exemplo, o Notebook ou arquivo tem o nome de agent.py. O Notebook ou arquivo deve conter um agente LangChain, denominado aqui como lc_agent.

  2. Inclua mlflow.models.set_model(lc_agent) no Notebook ou no arquivo.

  3. Crie outro notebook para servir como o notebook do driver (chamado driver.py neste exemplo).

  4. No Notebook do driver, use o código a seguir para executar agent.py e log os resultados para um modelo MLflow:

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

    O parâmetro resources declara Databricks-gerenciar recurso necessário para atender ao agente, como um índice de pesquisa de vetor ou atendimento endpoint que atende a um modelo de fundação. Para obter mais informações, consulte Autenticação para Databricks recurso.

  5. implantado o modelo. Veja implantado um agente para aplicativos generativos AI.

  6. Quando o ambiente de serviço for carregado, agent.py é executado.

  7. Quando chega uma solicitação de atendimento, lc_agent.invoke(...) é chamado.

Python

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 baseado em código com PyFunc

As instruções e o exemplo de código a seguir mostram ao senhor como log um agente com PyFunc.

  1. Crie um arquivo do Notebook ou Python com seu código. Neste exemplo, o Notebook ou arquivo tem o nome de agent.py. O Notebook ou arquivo deve conter uma classe PyFunc, denominada PyFuncClass.

  2. Incluir mlflow.models.set_model(PyFuncClass) no notebook ou arquivo.

  3. Crie outro notebook para servir como o notebook do driver (chamado driver.py neste exemplo).

  4. No Notebook do driver, use o código a seguir para executar agent.py e log os resultados para um modelo MLflow:

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

    O parâmetro resources declara Databricks-gerenciar recurso necessário para atender ao agente, como um índice de pesquisa de vetor ou atendimento endpoint que atende a um modelo de fundação. Para obter mais informações, consulte Autenticação para Databricks recurso.

  5. implantado o modelo. Veja implantado um agente para aplicativos generativos AI.

  6. Quando o ambiente de serviço for carregado, agent.py é executado.

  7. Quando chega uma solicitação de atendimento, PyFuncClass.predict(...) é chamado.

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

Autenticação para Databricks recurso

AI Os agentes geralmente precisam se autenticar em outro recurso para concluir a tarefa. Por exemplo, um agente pode precisar acessar um índice de pesquisa vetorial para consultar dados não estruturados.

Conforme descrito em Autenticação para recurso dependente, o servindo modelo suporta a autenticação tanto para Databricks-gerenciar quanto para recurso externo quando o senhor implanta o agente.

O servindo modelo suporta dois tipos diferentes de autenticação para Databricks-gerenciar recurso:

  1. Autenticação do sistema: Permite que a entidade de serviço do agente acesse qualquer recurso dependente especificado no momento do registro do agente. Isso é útil para acessar recursos compartilhados ou não confidenciais, por exemplo, um índice de pesquisa vetorial contendo documentação pública
  2. [Beta] Autenticação no pé do usuário: Permite que o agente use as credenciais do usuário final para acessar Databricks recurso. Isso é útil para cenários em que o agente precisa acessar dados confidenciais ou consultar APIs remotas para realizar ações por usuário

Especifique o recurso para passagem automática de autenticação (autenticação do sistema)

Para os tipos de recurso mais comuns do Databricks, o Databricks suporta e recomenda a declaração de dependências de recurso para o agente antecipadamente durante o registro. Isso permite a passagem automática de autenticação quando o senhor implanta o agente - Databricks provisiona, gira e gerencia automaticamente credenciais de curta duração para acessar com segurança essas dependências de recurso de dentro do agente endpoint.

Para ativar a passagem automática de autenticação, especifique o recurso dependente usando o parâmetro resources do logs() API, conforme mostrado no código a seguir.

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

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"),
DatabricksUCConnection(connection_name="your_connection_name"),
]
)

A Databricks recomenda que o senhor especifique manualmente resources para todos os tipos de agente.

nota

Se o senhor não especificar o recurso ao registrar os agentes do LangChain usando mlflow.langchain.log_model(...), o MLflow realizará a inferência automática de recurso com o melhor esforço possível. No entanto, isso pode não capturar todas as dependências, resultando em erros de autorização ao servir ou consultar o agente.

A tabela a seguir lista o recurso Databricks que oferece suporte à passagem automática de autenticação e a versão mínima mlflow necessária para log o recurso.

Tipo de recurso

Versão mínima mlflow necessária para log o recurso

Índice de pesquisa vetorial

Requer mlflow 2.13.1 ou acima

Modelo de ponto de extremidade de serviço

Requer mlflow 2.13.1 ou acima

Armazém SQL

Requer mlflow 2.16.1 ou acima

Função do Unity Catalog

Requer mlflow 2.16.1 ou acima

Espaço Genie

Requer mlflow 2.17.1 ou acima

Unity Catalog tabela

Requer mlflow 2.18.0 ou acima

Conexão do Unity Catalog

Requer mlflow 2.17.1 ou acima

Autenticação em nome do usuário

info

Beta

Esse recurso está na versão beta.

Ao registrar um agente que usa a autenticação no local do usuário, especifique o conjunto mínimo de escopos da API do Databricks necessários para executar ações como o usuário final no código do agente. Isso garante que o agente tenha acesso com privilégios mínimos para executar ações em nome do usuário final quando implantado, aumentando a segurança ao impedir ações não autorizadas e minimizando o risco de uso indevido de tokens.

abaixo é uma lista de escopos necessários para acessar vários tipos comuns de Databricks recurso.

Databricks recurso

Escopos de API necessários

Índice de pesquisa vetorial

serving.serving-endpoints, vectorsearch.vector-search-endpoints,vectorsearch.vector-search-indexes

Endpoint do serviço de modelos

serving.serving-endpoints

SQL Warehouse

sql.statement-execution,sql.warehouses

Conexões UC

catalog.connections

Espaço Genie

dashboards.genie

Para habilitar a autenticação na metade do usuário, passe um MLflow AuthPolicy para log_model(), conforme mostrado no exemplo abaixo. O MLflow AuthPolicy tem dois componentes:

  • system_auth_policy: Especifique o recurso para autenticação do sistema. Normalmente, os agentes usarão a autenticação do sistema para recursos compartilhados (por exemplo, para consultar o endpoint do modelo de serviço) em combinação com a autenticação no local do usuário para acessar recursos confidenciais ou APIs
  • user_auth_policy: Especifique os escopos de API de que seu agente precisa para autenticação no local do usuário
Python
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
)
)

Autenticação automática para clientes OpenAI

Se o seu agente usar o cliente OpenAI, use o SDK da Databricks para autenticar automaticamente durante a implantação. O SDK da Databricks fornece um wrapper para a construção do cliente OpenAI com autorização configurada automaticamente. Execute o seguinte em seu Notebook:

Python
% pip install databricks-sdk[openai]
Python
from databricks.sdk import WorkspaceClient
def openai_client(self):
w = WorkspaceClient()
return w.serving_endpoints.get_open_ai_client()

Em seguida, especifique o modelo de serviço endpoint como parte de resources para autenticar automaticamente no momento da implementação.

registrar o agente para Unity Catalog

Antes de implantar o agente, o senhor deve registrar o agente em Unity Catalog. registro do agente pacote como um modelo em Unity Catalog. Como resultado, o senhor pode usar as permissões Unity Catalog para autorização de recurso no agente.

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

Próximas etapas