Registrar e registro agentes de AI (Model Serving)

info

Para novos casos de uso, a Databricks recomenda implantar agentes em Databricks Apps para controle total sobre o código do agente, a configuração do servidor e o fluxo de trabalho de implantação. Consulte Criar um agente de AI e implantá-lo em Databricks Apps. Para migrar um agente existente, consulte Migrar um agente de Model Serving para Databricks Apps.

Registre agentes de AI usando Agentes Personalizados. Registrar um agente forma a base do processo de desenvolvimento. O registro em log captura um "ponto no tempo" do código e da configuração do agente para que você possa avaliar a qualidade dele.

Requisitos

Crie um agente de AI antes de registrá-lo.

A Databricks recomenda instalar a versão mais recente do databricks-sdk.

Python

% pip install databricks-sdk

Log baseado em código

A Databricks recomenda usar a funcionalidade Modelos de Código 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 pacotes. Quando o agente é implantado, o ambiente Python é restaurado e o código do agente é executado para carregar o agente na memória, para que possa ser invocado quando o endpoint for chamado.

Você pode combinar esta abordagem com o uso de APIs de validação pré-implantação como mlflow.models.predict() para garantir que o agente seja executado de forma confiável quando implantado para disponibilização.

Para ver um exemplo de registro baseado em código, consulte Notebooks de exemplo de autoria do ResponsesAgent.

Inferir Assinatura do Modelo durante o registro

nota

A Databricks recomenda desenvolver um agente usando a interface ResponsesAgent. Se estiver usando ResponsesAgent, você pode pular esta seção; o MLflow infere automaticamente uma assinatura válida para seu agente.

Se não estiver usando a interface ResponsesAgent, você deve usar um dos seguintes métodos para especificar a Assinatura do modelo MLflow do seu agente no momento do registro:

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 que você fornece. Essa abordagem é mais conveniente do que definir manualmente a assinatura.

A assinatura do modelo MLflow valida entradas e saídas para garantir que o agente interaja corretamente com o sistema. ferramentas subsequentes como o AI Playground e o aplicativo de avaliação. Também orienta outros aplicativos sobre como usar o agente. eficazmente.

Os exemplos de LangChain e PyFunc abaixo usam a inferência de Model Signature.

Se você preferir definir explicitamente uma assinatura de modelo no momento do registro, consulte Documentação do MLflow - Como registrar modelos com assinaturas.

Registro baseado em código com LangChain

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

1. Crie um notebook ou arquivo Python com seu código. Para este exemplo, o notebook ou arquivo é denominado agent.py. O notebook ou arquivo deve conter um agente LangChain, referido aqui como lc_agent.

2. Inclua mlflow.models.set_model(lc_agent) 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 seguinte código para executar agent.py e registrar os resultados em um modelo MLflow:

   Python

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

   O parâmetro resources declara recursos gerenciados pelo Databricks necessários para servir o agente, como um índice de AI Search ou um endpoint de disponibilização que disponibiliza um modelo de base. Para obter mais informações, consulte Implementar o passthrough de autenticação automático.

5. Implante o modelo. Consulte Implante um agente para aplicativos de AI generativa (Model Serving).

6. Quando o ambiente de serviço é 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 seguintes instruções e amostras de código mostram como registrar um agente com PyFunc.

1. Crie um notebook ou arquivo Python com seu código. Para este exemplo, o notebook ou arquivo é denominado agent.py. O notebook ou arquivo deve conter uma classe PyFunc, nomeada PyFuncClass.

2. Inclua 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 seguinte código para executar agent.py e use log_model() para registrar os resultados em um modelo MLflow:

   Python

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

   O parâmetro resources declara recursos gerenciados pelo Databricks necessários para servir o agente, como um índice de AI Search ou um endpoint de disponibilização que disponibiliza um modelo de base. Para obter mais informações, consulte Implementar o passthrough de autenticação automático.

5. Implante o modelo. Consulte Implante um agente para aplicativos de AI generativa (Model Serving).

6. Quando o ambiente de serviço é 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-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)

Autenticação para recursos do Databricks

Agentes de AI frequentemente precisam se autenticar a outros recursos para concluir tarefas. Por exemplo, um agente implantado pode precisar acessar um índice de Pesquisa de AI para query dados não estruturados ou acessar o Registro de Prompts para carregar prompts dinâmicos.

Passagem de autenticação automática e autenticação em nome de requerem configuração durante o registro do agente.

Registre o agente no Unity Catalog

Antes de implantar o agente, você deve registrá-lo no Unity Catalog. O registro do agente o empacota como um modelo no Unity Catalog. Como resultado, você pode usar as permissões do Unity Catalog para autorização dos recursos 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)

Consulte mlflow.register_model().

Outros recursos

• Adicione rastreamentos a um agente de AI.
• Implantar um agente de AI.