Registre em log e registre agentes de IA

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.

Registrar em log com base em código versus serialização

Você pode usar o registro em log do MLflow baseado em código ou baseado em serialização. O Databricks recomenda que você use o registro em log com base em código.

Registro em log do MLflow com base em código: o código da cadeia é capturado como um arquivo do Python. O ambiente do Python é capturado como uma lista de pacotes. Quando a cadeia é implantada, o ambiente do Python é restaurado e o código da cadeia é executado a fim de carregar a cadeia na memória para que ela possa ser invocada quando o ponto de extremidade for chamado.

Registro em log do MLflow com base em serialização: o código da cadeia e o estado atual no ambiente do Python são serializados em disco, muitas vezes usando bibliotecas como pickle ou joblib. Quando a cadeia é implantada, o ambiente do Python é restaurado e o objeto serializado é carregado na memória para que possa ser invocado quando o ponto de extremidade é chamado.

A tabela mostra as vantagens e desvantagens de cada método.

Para o registro em log com base em código, o código que registra seu agente ou cadeia deve estar em um notebook separado do código da cadeia. Este notebook é chamado de notebook de driver. Para visualizar um notebook de exemplo, consulte Notebooks de exemplo.

Registro em log com base em código com LangChain

1. Crie um notebook ou arquivo Python com seu código. Para fins deste exemplo, o notebook ou arquivo é nomeado chain.py. O notebook ou arquivo deve conter uma cadeia LangChain, conhecida aqui como lc_chain.
2. Inclua mlflow.models.set_model(lc_chain) 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 mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”) para executar chain.py e registrar os resultados em um modelo do MLflow.
5. Implantar o modelo. Confira Implantar um agente para um aplicativo de IA generativa. A implantação do agente pode depender de outros recursos do Databricks, como um índice de pesquisa de vetor e pontos de extremidade que servem a modelos. Para agentes do LangChain:
   • O MLflow log_model infere as dependências exigidas pela cadeia e registra-as no arquivo MLmodel no artefato do modelo registrado. A partir do Mlflow versão 2.17.0, você pode substituir essas dependências inferidas. Consulte Especificar recursos para o agente PyFunc ou LangChain.
   • Durante a implantação, databricks.agents.deploy cria automaticamente os tokens OAuth M2M necessários para acessar e se comunicar com essas dependências de recursos inferidas.
6. Quando o ambiente de serviço é carregado, o chain.py é executado.
7. Quando uma solicitação de serviço entra, o lc_chain.invoke(...) é chamado.

import mlflow

code_path = "/Workspace/Users/first.last/chain.py"
config_path = "/Workspace/Users/first.last/config.yml"

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using LangChain
with mlflow.start_run():
  logged_chain_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this is the configuration that is used for training the model. The development_config is overwritten.
    artifact_path="chain", # 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 your chain
    example_no_conversion=True, # Required
  )

print(f"MLflow Run: {logged_chain_info.run_id}")
print(f"Model URI: {logged_chain_info.model_uri}")

# To verify that the model has been logged correctly, load the chain and call `invoke`:
model = mlflow.langchain.load_model(logged_chain_info.model_uri)
model.invoke(example)

Registro em log com base em código com PyFunc

1. Crie um notebook ou arquivo Python com seu código. Para fins deste exemplo, o notebook ou arquivo é nomeado chain.py. O notebook ou arquivo deve conter uma classe PyFunc, conhecida aqui como 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 mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”, resources=”/path/to/resources.yaml”) para executar chain.py e registrar os resultados em um modelo do MLflow. O resources parâmetro declara todos os recursos necessários para atender ao modelo, como um índice de pesquisa vetorial ou um ponto de extremidade de serviço que atende a um modelo de base. Para obter um exemplo de arquivo de recursos para PyFunc, consulte Especificar recursos para o agente PyFunc ou LangChain.
5. Implantar o modelo. Confira Implantar um agente para um aplicativo de IA generativa.
6. Quando o ambiente de serviço é carregado, o chain.py é executado.
7. Quando uma solicitação de serviço entra, o PyFuncClass.predict(...) é chamado.

import mlflow

code_path = "/Workspace/Users/first.last/chain.py"
config_path = "/Workspace/Users/first.last/config.yml"

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using PyFunc model

resources_path = "/Workspace/Users/first.last/resources.yml"

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
  )

print(f"MLflow Run: {logged_chain_info.run_id}")
print(f"Model URI: {logged_chain_info.model_uri}")

# To verify that the model has been logged correctly, load the chain and call `invoke`:
model = mlflow.pyfunc.load_model(logged_chain_info.model_uri)
model.invoke(example)

Especificar recursos para o agente PyFunc ou LangChain

Você pode especificar recursos, como um índice de busca em vetores e um ponto de extremidade de serviço, que são necessários para servir o modelo.

Para o LangChain, os recursos são detectados e registrados automaticamente com o modelo usando uma abordagem de melhor esforço. A partir do MLflow versão 2.17.0, você pode substituir esses recursos inferidos automaticamente usando um código semelhante ao mostrado abaixo. Isso é recomendado para casos de uso de produção, pois permite garantir que os agentes sejam registrados com as dependências necessárias.

Ao implantar um agente com sabor pyfunc, você precisa adicionar manualmente quaisquer dependências de recurso do agente implantado. Um token OAuth M2M com acesso a todos os recursos especificados no parâmetro resources é criado e fornecido ao agente implantado.

O código a seguir especifica dependências usando o resources parâmetro.

import mlflow
from mlflow.models.resources import (
    DatabricksFunction,
    DatabricksServingEndpoint,
    DatabricksSQLWarehouse,
    DatabricksVectorSearchIndex,
)

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
      DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
      DatabricksFunction(function_name="ml.tools.python_exec"),
    ]
  )

Você também pode adicionar recursos especificando-os em um arquivo resources.yaml. Você pode referenciar esse caminho de arquivo no parâmetro resources. Um token OAuth M2M com acesso a todos os recursos especificados no resources.yaml é criado e fornecido ao agente implantado.

Veja a seguir um arquivo de exemplo resources.yaml que define o modelo que atende pontos de extremidade e um índice de pesquisa de vetor.

api_version: "1"
databricks:
  vector_search_index:
    - name: "catalog.schema.my_vs_index"
  serving_endpoint:
    - name: databricks-dbrx-instruct
    - name: databricks-bge-large-en

Registre a cadeia no Catálogo do Unity

Antes de implantar a cadeia, registre a cadeia no Catálogo do Unity. Ao registrar a cadeia, ela é empacotada como um modelo no Catálogo do Unity e você poderá usar permissões do Catálogo do Unity para autorização dos recursos na cadeia.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "chain_name"

model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_chain_info.model_uri, name=model_name)

Próximas etapas

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