Log e registro de agentes de IA

Prévia

Registrar agentes de IA usando Mosaic AI a estrutura de agentes . 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 de IA antes de registrá-lo.

Log baseado em código versus baseado em serialização

O senhor pode usar o registro do MLflow baseado em código ou o registro do MLflow baseado em serialização. A Databricks recomenda que o senhor use o registro baseado em código.

Registro de MLflow baseado em código: O código da cadeia é capturado como um arquivo Python. O ambiente Python é capturado como uma lista de pacotes. Quando a cadeia é implantada, o ambiente Python é restaurado e o código da cadeia é executado para carregar a cadeia na memória, de modo que possa ser invocada quando o endpoint for chamado.

Registro em MLflow baseado em serialização: O código da cadeia e o estado atual no ambiente Python são serializados para o disco, geralmente usando biblioteca como pickle ou joblib. Quando a cadeia é implantada, o ambiente Python é restaurado e o objeto serializado é carregado na memória para que possa ser invocado quando o endpoint for chamado.

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

Método	Vantagens	Desvantagens
Log de MLflow baseado em código	Supera as limitações inerentes à serialização, que não é suportada por muitas bibliotecas populares do GenAI. Salva uma cópia do código original para referência posterior. Não há necessidade de reestruturar seu código em um único objeto que possa ser serializado.	`log_model(...)` deve ser chamado de um notebook diferente do código da cadeia (chamado de notebook do driver).
Registro de MLflow baseado em serialização	`log_model(...)` pode ser chamado no mesmo notebook em que o modelo é definido.	O código original não está disponível. Todas as bibliotecas e objetos usados na cadeia devem ser compatíveis com a serialização.

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

Registro de log baseado em código com LangChain

Crie um notebook ou arquivo Python com seu código. Para fins deste exemplo, o Notebook ou arquivo é denominado chain.py. O Notebook ou arquivo deve conter uma cadeia LangChain, referenciada aqui como lc_chain.
Incluir mlflow.models.set_model(lc_chain) no notebook ou arquivo.
Crie outro notebook para servir como o notebook do driver (chamado driver.py neste exemplo).
No Notebook do driver, use mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”)para executar chain.py e log os resultados para um modelo MLflow.
implantado o modelo. Veja implantado um agente para aplicação de IA generativa. A implementação do seu agente pode depender de outros Databricks recursos, como um índice de pesquisa de vetores e um endpoint de modelo de serviço. Para agentes da LangChain:
- O MLflow log_model infere as dependências exigidas pela cadeia e as logs para o arquivo MLmodel no artefato de modelos registrados.
- Durante a implementação, databricks.agents.deploy cria automaticamente o M2M OAuth tokens necessário para acessar e se comunicar com essas dependências de recursos inferidas.
Quando o ambiente de serviço for carregado, chain.py é executado.
Quando chega uma solicitação de atendimento, 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 baseado em código com PyFunc

Crie um notebook ou arquivo Python com seu código. Para fins deste exemplo, o Notebook ou arquivo é denominado chain.py. O notebook ou arquivo deve conter uma classe PyFunc, referenciada aqui como PyFuncClass.
Incluir mlflow.models.set_model(PyFuncClass) no notebook ou arquivo.
Crie outro notebook para servir como o notebook do driver (chamado driver.py neste exemplo).
No Notebook do driver, use mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”, resources=”/path/to/resources.yaml”) para executar chain.py e log os resultados para um modelo MLflow. O parâmetro resources declara qualquer recurso necessário para atender ao modelo, como, por exemplo, um índice de pesquisa vetorial ou o serviço endpoint que atende a um modelo de fundação. Veja um exemplo de arquivo de recurso para o PyFunc.
implantado o modelo. Veja implantado um agente para aplicação de IA generativa.
Quando o ambiente de serviço for carregado, chain.py é executado.
Quando chega uma solicitação de atendimento, 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 o recurso para o agente PyFunc

O senhor pode especificar recursos, como um índice de pesquisa vetorial e um serviço endpoint, que são necessários para servir o modelo. Para LangChain, os recursos são automaticamente coletados e registrados junto com o modelo.

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

Observação

O senhor pode substituir o recurso ao qual o site endpoint tem permissão especificando manualmente o recurso ao registrar a cadeia.

A seguir, mostramos como adicionar dependências de endpoint de atendimento e de índice de pesquisa de vetores especificando-as no parâmetro resources.

 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="rag.studio_bugbash.databricks_docs_index")
        ]
   )

O senhor 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 M2M OAuth com acesso a todos os recursos especificados no resources.yaml é criado e fornecido ao agente implantado.

A seguir, um exemplo de arquivo resources.yaml que define o endpoint do modelo de serviço e um índice de pesquisa de vetores.

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 Unity Catalog

Antes de implantar a cadeia, você deve registrar a cadeia no Unity Catalog. Quando o usuário registra a cadeia, ela é empacotada como um modelo no Unity Catalog, e o usuário pode utilizar as permissões do Unity Catalog para autorização de 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)