Registrar e registrar os agentes de IA

Registrar agentes de IA usando o Mosaic AI Agent Framework. 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.

Requirements

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

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

% pip install databricks-sdk

Registro em log baseado em código

O Databricks recomenda usar os Modelos de funcionalidade do Código do MLflow ao registrar agentes 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.

Para ver um exemplo de log baseado em código, consulte blocos de anotações de exemplo de criação responsesAgent.

Inferir a assinatura de modelo durante o registro em log

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

1. Definir 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 do MLflow 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.

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 Documentos do 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 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 Implementar passagem de autenticação automática.

5. Implantar o modelo. Consulte Implantar um agente para aplicativos de IA generativos.

6. Quando o ambiente de serviço é carregado, 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 usar log_model() para registrar os resultados em um modelo de 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 de vetor ou um endpoint de serviço que serve a um modelo de base. Para obter mais informações, consulte Implementar passagem de autenticação automática.

5. Implantar o modelo. Consulte Implantar um agente para aplicativos de IA generativos.

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

Os agentes de IA geralmente precisam se autenticar em outros recursos para concluir tarefas. Por exemplo, um agente implantado pode precisar acessar um índice de Pesquisa de Vetor para consultar dados não estruturados ou acessar o Registro de Prompt para carregar prompts dinâmicos.

A passagem de autenticação automática e a autenticação em nome de outrem necessitam de configuração durante o log do agente.

Registrar o agente no Catálogo do Unity

Antes de implantar o agente, você deve registrar o agente no Catálogo do Unity. O registro do 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)

Consulte mlflow.register_model().

Próximas etapas

• Adicionar rastreamentos a um agente de IA.
• Implantar um agente de IA.