Registre em log e registre agentes de IA

Importante

Esse recurso está em uma versão prévia.

Registre agentes de IA usando a Estrutura de Agente de IA do Mosaico. Registrar um agente em log é 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 da configuração.

Requisitos

Crie um agente de IA antes de registrá-lo em log.

Registro em log baseado em código

O Databricks recomenda usar os modelos da funcionalidade Modelos de Código do MLflow ao registrar agentes de registro em log.

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

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

O agente ou o código que registra o agente precisa estar em um notebook separado do código do agente. Este notebook é chamado de notebook de driver. Para visualizar um notebook de exemplo, consulte Notebooks de exemplo.

Inferir a Assinatura de Modelo durante o registro em log

Durante o registro em log, você precisa definir uma Assinatura de Modelo do MLflow, que especifica o esquema de entrada e saída do agente. A assinatura valida entradas e saídas para garantir que o agente interaja corretamente com ferramentas downstream como o AI Playground e o aplicativo de revisão. Ele também orienta outros aplicativos sobre como usar o agente efetivamente.

O Databricks recomenda usar 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.

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

Se você preferir definir explicitamente uma Assinatura de Modelo por conta própria no momento do registro em log, consulte Documentação do MLflow – Como registrar modelos com assinaturas.

Registro em log com base em código com LangChain

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

1. Crie um notebook ou arquivo Python com seu código. Para este exemplo, o bloco de anotações ou arquivo é nomeado agent.py. O bloco de anotações ou arquivo deve conter um agente LangChain, conhecido aqui como lc_agent.

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

3. Crie um notebook para servir como o notebook de driver (chamado driver.py neste exemplo).

4. No notebook de driver, use o seguinte código para executar agent.py e registrar os resultados em um modelo MLflow:

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

   O parâmetro resources declara os recursos gerenciados pelo Databricks necessários para servir ao agente, como um índice de pesquisa de vetor ou um endpoint de serviço que serve a um modelo de base. Para obter mais informações, consulte Como especificar recursos para passagem automática de autenticação.

5. Implantar o modelo. Confira Implantar um agente para um aplicativo de IA generativa.

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

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

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)

Log baseado em código com PyFunc

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

1. Crie um notebook ou arquivo Python com seu código. Para este exemplo, o bloco de anotações ou arquivo é nomeado agent.py. O bloco de anotações ou arquivo deve conter uma classe PyFunc, chamada PyFuncClass.

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

3. Crie um notebook para servir como o notebook de driver (chamado driver.py neste exemplo).

4. No notebook de driver, use o seguinte código para executar agent.py e registrar os resultados em um modelo MLflow:

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

   O parâmetro resources declara os recursos gerenciados pelo Databricks necessários para atender ao agente, como um índice de busca em vetores ou um ponto de extremidade do serviço que disponibiliza um modelo de base. Para obter mais informações, consulte Como especificar recursos para passagem automática de autenticação.

5. Implantar o modelo. Confira Implantar um agente para um aplicativo de IA generativa.

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

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

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)

Especificar recursos para a passagem de autenticação automática

Os agentes de IA geralmente precisam se autenticar em outros recursos para concluir tarefas. Por exemplo, um agente pode precisar acessar um índice de Pesquisa de Vetor para consultar dados não estruturados.

Conforme descrito em Autenticação para recursos dependentes, o Model Serving dá suporte à autenticação para recursos gerenciados pelo Databricks e externos ao implantar o agente.

Para os tipos de recursos mais comuns do Databricks, o Databricks dá suporte e recomenda declarar dependências de recursos para o agente previamente durante o registro em log. Isso permite a passagem de autenticação automática quando você implanta o agente. O Databricks provisiona, gira e gerencia automaticamente credenciais de curta duração para acessar com segurança essas dependências de recursos de dentro do ponto de extremidade do agente.

Para habilitar a passagem de autenticação automática, especifique os recursos dependentes usando o parâmetro resources da API log_model(), conforme mostrado no código a seguir.

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

O Databricks recomenda que você especifique manualmente resources para todos os tipos de agente.

Observação

Se você não especificar recursos ao registrar agentes LangChain usando mlflow.langchain.log_model(...), o MLflow executará a inferência automática de recursos da melhor maneira 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 os recursos do Databricks que dão suporte à passagem de autenticação automática e à versão mínima de mlflow necessária para registrar o recurso em log.

Tipo de recurso	A versão mínima de mlflow é exigida para registrar o recurso em log
Índice de Busca Vetorial	Requer mlflow 2.13.1 ou superior
Ponto de extremidade do Serviço de Modelo	Requer mlflow 2.13.1 ou superior
SQL Warehouse	Requer mlflow 2.16.1 ou superior
Função do Unity Catalog	Requer mlflow 2.16.1 ou superior
Espaço do gênio	Requer mlflow 2.17.1 ou superior
Tabela do Catálogo Unity	Requer mlflow 2.18.0 ou superior

Registrar o agente no Catálogo do Unity

Antes de implantar o agente, você deve registrar o agente no Catálogo do Unity. Registrar o agente o empacota como um modelo no Catálogo do Unity. Como resultado, você pode usar permissões do Catálogo do Unity para autorização para recursos no 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)

Próximas etapas

• Adicione rastreamentos a um agente de IA.
• Implante um agente de IA.