Registar e inscrever agentes de IA

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 "momento específico" do código e da configuração do agente para que o utilizador possa avaliar a qualidade da configuração.

Requirements

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

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

% pip install databricks-sdk

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 lista de pacotes. Quando o agente é implantado, o ambiente Python é reestabelecido e o código do agente é executado para carregar o agente na memória para que ele possa ser invocado quando o endpoint é chamado.

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.

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

Inferir assinatura do modelo durante o registo

Se não usar a interface ResponsesAgent, deverá usar um dos seguintes métodos para especificar a assinatura de modelo MLflow do seu agente na altura do registo:

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 fornecido. 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 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.

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 notas para servir como o bloco de notas principal (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 geridos pelo Databricks necessários para o agente, como um índice de pesquisa vetorial ou um endpoint que serve um modelo base. Para obter mais informações, consulte Implementar passagem de autenticação automática.

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

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

7. Quando chega um pedido de atendimento, lc_agent.invoke(...) será ativado.

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 notas para servir como o bloco de notas principal (chamado driver.py neste exemplo).

4. No bloco de anotações do driver, use o seguinte código para executar agent.py e usar log_model() para 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 geridos pelo Databricks necessários para o agente, como um índice de pesquisa vetorial ou um endpoint que serve um modelo base. Para obter mais informações, consulte Implementar passagem de autenticação automática.

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

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

7. Quando chega um pedido de atendimento, PyFuncClass.predict(...) será ativado.

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 Vetorial 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 exigem configuração durante o registro em log do agente.

Registrar o agente no Catálogo Unity

Antes de implantar o agente, você deve registrá-lo no Unity Catalog. Ao registar o agente, empacota-o como um modelo no Unity Catalog. Como resultado, você pode usar as permissões do Catálogo Unity 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)

Consulte mlflow.register_model().

Próximos passos

• Adicione traços a um agente de IA.
• Implante um agente de IA.