Registrar e registrar agentes de IA

Importante

Esta funcionalidade está em Pré-visualização Pública.

Registre agentes de IA usando o Mosaic AI Agent Framework. O registro de um agente é a base do processo de desenvolvimento. O registro em log captura um "point-in-time" 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.

Registo baseado em código

O 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 list de pacotes. Quando o agente é implantado, o ambiente Python é restaurado. O código do agente é então executado para carregar o agente na memória, permitindo que ele seja invocado quando se chama o endpoint.

Você pode combinar essa 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 servir.

O código que registra o agente ou agente deve estar em um bloco de anotações separado do código do agente. Este notebook é chamado de notebook de driver. Para obter um bloco de anotações de exemplo, consulte Blocos de anotações de exemplo.

Inferir assinatura de modelo durante o registro em log

Durante o registro, você deve definir um MLflow Model Signature, que especifica a entrada e saída do agente schema. 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 de forma eficaz.

O Databricks recomenda o uso dos recursos de inferência de assinatura de modelo do MLflow para generate automaticamente a assinatura do agente com base em um exemplo de entrada fornecido. Essa abordagem é mais conveniente do que definir manualmente a assinatura.

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

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

Registro em log baseado em código com LangChain

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

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

2. Inclua mlflow.models.set_model(lc_agent) no bloco de notas ou ficheiro.

3. Crie um novo bloco de anotações para servir como o bloco de anotações de driver (chamado driver.py neste exemplo).

4. No bloco de anotações do driver, use o código a seguir 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 vetorial ou um ponto de extremidade de serviço que serve a um modelo de base. Para obter mais informações, consulte Especificar recursos para passagem automática de autenticação.

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

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

7. Quando um pedido de serviço chega, 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)

Registro baseado em código com PyFunc

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

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

2. Inclua mlflow.models.set_model(PyFuncClass) no bloco de anotações ou arquivo.

3. Crie um novo bloco de anotações para servir como o bloco de anotações de driver (chamado driver.py neste exemplo).

4. No bloco de anotações do driver, use o código a seguir 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 servir ao agente, como um índice de pesquisa vetorial ou um ponto de extremidade de serviço que serve a um modelo de base. Para obter mais informações, consulte Especificar recursos para passagem automática de autenticação.

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

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

7. Quando um pedido de serviço chega, 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 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 Vetorial para consultar dados não estruturados.

Conforme descrito em Autenticação para recursos dependentes, o Model Serving oferece suporte à autenticação em recursos gerenciados e externos pelo Databricks quando você implanta o agente.

Para os tipos de recursos Databricks mais comuns, o Databricks suporta e recomenda declarar dependências de recursos para o agente antecipadamente durante o registro. Isso permite de passagem de autenticação automática quando você implanta o agente - o Databricks provisiona, gira e gerencia automaticamente credentials 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 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.

Nota

Se não especificar recursos ao registar agentes LangChain usando mlflow.langchain.log_model(...), o MLflow realizará automaticamente a inferência de recursos da melhor forma 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 seguinte table lista os recursos do Databricks que oferecem suporte à passagem de autenticação automática e a versão mínima de mlflow necessária para registar o recurso.

Tipo de recurso	Versão mínima de mlflow necessária para registrar o recurso
Índice de pesquisa vetorial	Requer mlflow 2.13.1 ou superior
Ponto de extremidade de serviço de modelo	Requer mlflow 2.13.1 ou superior
Armazém SQL	Requer mlflow 2.16.1 ou superior
Função da Unidade Catalog	Requer mlflow 2.16.1 ou superior
Espaço Genie	Requer mlflow 2.17.1 ou superior
Unidade Catalogtable	Requer mlflow 2.18.0 ou superior

Registre o agente no Unity Catalog

Antes de implantar o agente, você deve registrar o agente no Unity Catalog. Ao registar o agente, empacota-o como um modelo no Unity Catalog. Como resultado, você pode usar as permissões do Unity Catalog para autorização de 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óximos passos

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