registro e registro AI agentes

Prévia

Esse recurso está em Pré-lançamento público.

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.

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.

O código que logs o agente ou agente deve estar em um Notebook separado do código do agente. Esse notebook é chamado de notebook de driver. Para ver um exemplo de Notebook, consulte Example Notebook.

Inferir a assinatura do modelo durante o registro

Durante o registro, o senhor deve 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 posteriores, como o AI Playground e o aplicativo de avaliação. Ele também orienta outros aplicativos sobre como usar o agente de forma eficaz.

A 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 fornecido pelo senhor. 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 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:

    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 vetores ou o serviço endpoint que atende a um modelo de fundação. Para obter mais informações, consulte Especificar recurso para passagem automática de autenticação.

  5. implantado o modelo. Veja implantado um agente para aplicação generativa do 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.


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:

    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 vetores ou o serviço endpoint que atende a um modelo de fundação. Para obter mais informações, consulte Especificar recurso para passagem automática de autenticação.

  5. implantado o modelo. Veja implantado um agente para aplicação generativa do AI .

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

  7. Quando chega uma solicitação de atendimento, 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)

Especifique o recurso para a passagem automática de autenticação

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.

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.

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

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

Observação

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 suporta a 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

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.

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)