Crie e registre agentes de IA

Prévia

Este artigo mostra como criar e registrar agentes de IA, como aplicações RAG, utilizando o Mosaic AI Agent Framework.

O que são cadeias e agentes?

Os sistemas de IA geralmente têm muitos componentes. Por exemplo, um sistema de IA pode recuperar documentos de um índice vetorial, utilizar esses documentos para complementar o texto de prompt e utilizar um modelo de base para resumir a resposta. O código que vincula esses componentes, também chamados de etapas, é chamado de cadeia.

Um agente é um sistema de IA muito mais avançado que se baseia em grandes modelos de linguagem para tomar decisões sobre os passos a serem seguidos com base na entrada. Em contraste, as cadeias são sequências codificadas de passos destinadas a alcançar um resultado específico.

Com o Agent Framework, você pode utilizar quaisquer bibliotecas ou pacotes para criar código. O Agent Framework também facilita a iteração do código enquanto você o desenvolve e testa. Você pode preparar arquivos de configuração que possibilitam a alteração dos parâmetros de código de forma rastreável sem precisar modificar o código real.

Requisitos

Para agentes que usam um índice de pesquisa vetorial gerenciado pelo Databricks, a versão mlflow 2.13.1 ou superior é necessária para usar a autorização automática com o índice vetorial.

Esquema de entrada do agente RAG

Os seguintes formatos de entrada são compatíveis com sua cadeia.

(Recomendado) Consultas utilizando o esquema de conclusão de bate-papo OpenAI. Deve ter uma matriz de objetos como parâmetro messages . Este formato é melhor para aplicações RAG.

question = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-Augmented Generation?",
        },
        {
            "role": "assistant",
            "content": "RAG, or Retrieval Augmented Generation, is a generative AI design pattern that combines a large language model (LLM) with external knowledge retrieval. This approach allows for real-time data connection to generative AI applications, improving their accuracy and quality by providing context from your data to the LLM during inference. Databricks offers integrated tools that support various RAG scenarios, such as unstructured data, structured data, tools & function calling, and agents.",
        },
        {
            "role": "user",
            "content": "How to build RAG for unstructured data",
        },
    ]
}

SplitChatMessagesRequest. Recomendado para aplicações de bate-papo com vários turnos, especialmente quando o usuário deseja gerenciar a consulta atual e o histórico separadamente.

{
"query": "What is MLflow",
"history": [
  {
  "role": "user",
  "content": "What is Retrieval-augmented Generation?"
  },
  {
  "role": "assistant",
  "content": "RAG is"
  }
  ]
}

Para a LangChain, a Databricks recomenda escrever sua cadeia em LangChain Expression Language. Em seu código de definição de cadeia, você pode utilizar um itemgetter para receber as mensagens ou os objetos query ou history, dependendo do formato de entrada que estiver utilizando.

Esquema de saída para o agente RAG

Seu código deve estar em conformidade com um dos seguintes formatos de saída suportados:

(Recomendado) ChatCompletionResponse. Esse formato é recomendado para clientes com interoperabilidade de formato de resposta OpenAI.
StringResponse. Este formato é o mais fácil e simples de interpretar.

Para LangChain, use StrOutputParser() como o passo final da cadeia. Sua saída deve retornar um único valor de string.

  chain = (
      {
          "user_query": itemgetter("messages")
          | RunnableLambda(extract_user_query_string),
          "chat_history": itemgetter("messages") | RunnableLambda(extract_chat_history),
      }
      | RunnableLambda(fake_model)
      | StrOutputParser()
  )

Se você estiver utilizando PyFunc, a Databricks recomenda utilizar dicas de tipo para anotar a função predict() com classes de dados de entrada e saída que são subclasses de classes definidas em mlflow.models.rag_signatures.

Você pode construir um objeto de saída a partir da classe de dados dentro de predict() para garantir que o formato seja seguido. O objeto retornado deve ser transformado em uma representação de dicionário para garantir que possa ser serializado.

Use parâmetros para controlar a iteração de qualidade

Na Agent Framework, você pode utilizar parâmetros para controlar como os agentes são executados. Isso possibilita que você faça uma iteração rápida variando as características do agente sem alterar o código. Os parâmetros são pares chave-valor que você define em um dicionário no Python ou em um arquivo .yaml.

Para configurar o código, você cria um ModelConfig, um conjunto de parâmetros de chave-valor. O ModelConfig é um dicionário no Python ou um arquivo .yaml . Por exemplo, você pode utilizar um dicionário durante o desenvolvimento e, em seguida, convertê-lo em um arquivo .yaml para implantação de produção e CI/CD. Para obter detalhes sobre ModelConfig, consulte a documentação do MLflow.

Um exemplo ModelConfig é mostrado abaixo.

llm_parameters:
  max_tokens: 500
  temperature: 0.01
model_serving_endpoint: databricks-dbrx-instruct
vector_search_index: ml.docs.databricks_docs_index
prompt_template: 'You are a hello world bot. Respond with a reply to the user''s
  question that indicates your prompt template came from a YAML file. Your response
  must use the word "YAML" somewhere. User''s question: {question}'
prompt_template_input_vars:
- question

Para chamar a configuração do seu código, use uma das seguintes opções:

# Example for loading from a .yml file
config_file = "configs/hello_world_config.yml"
model_config = mlflow.models.ModelConfig(development_config=config_file)

# Example of using a dictionary
config_dict = {
    "prompt_template": "You are a hello world bot. Respond with a reply to the user's question that is fun and interesting to the user. User's question: {question}",
    "prompt_template_input_vars": ["question"],
    "model_serving_endpoint": "databricks-dbrx-instruct",
    "llm_parameters": {"temperature": 0.01, "max_tokens": 500},
}

model_config = mlflow.models.ModelConfig(development_config=config_dict)

# Use model_config.get() to retrieve a parameter value
value = model_config.get('sample_param')

Registre o agente

O registro em log de um agente é a base do processo de desenvolvimento. O registro em log captura um "point-in-time" do código e da configuração do agente para você avaliar a qualidade da configuração. Ao desenvolver agentes, o Databricks recomenda que você use o log baseado em código em vez do log baseado em serialização. Para obter mais informações sobre os prós e contras de cada tipo de log, consulte Log baseado em código versus baseado em serialização.

Esta seção aborda como utilizar o registro em log baseado em código. Para obter detalhes sobre como utilizar o registro baseado em serialização, consulte Fluxo de trabalho de registro baseado em serialização.

Fluxo de trabalho de registro baseado em código

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.

Fluxo de trabalho de registro 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, inclua a chamada mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”). Esta chamada executa chain.py e registra os resultados em um modelo MLflow.
Implante o modelo.
Quando o ambiente de serviço for carregado, chain.py é executado.
Quando chega uma solicitação de atendimento, lc_chain.invoke(...) é chamado.

Fluxo de trabalho de 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, inclua a chamada mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”). Esta chamada executa chain.py e registra os resultados em um modelo MLflow.
Implante o modelo.
Quando o ambiente de serviço for carregado, chain.py é executado.
Quando chega uma solicitação de atendimento, PyFuncClass.predict(...) é chamado.

Código de exemplo para cadeias de log

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
  )

# or use a PyFunc model
# 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,
#   )

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

Para verificar se o modelo foi registrado corretamente, carregue a cadeia e chame invoke:

# Using LangChain
model = mlflow.langchain.load_model(logged_chain_info.model_uri)
model.invoke(example)

# Using PyFunc
model = mlflow.pyfunc.load_model(logged_chain_info.model_uri)
model.invoke(example)

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)

Notebooks de Exemplo

Esses notebooks criam uma cadeia simples "Hello, world" para ilustrar como criar uma aplicação de cadeia no Databricks. O primeiro exemplo cria uma cadeia simples. O segundo notebook de exemplo ilustra como utilizar parâmetros para minimizar as alterações de código durante o desenvolvimento.

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

Para criar e logar uma cadeia, você pode utilizar o log MLflow baseado em código ou o log MLflow baseado em serialização. A Databricks recomenda que você use o log baseado em código.

Com o log 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 é chamado.

Com o registro de MLflow baseado em serialização, o código da cadeia e o estado atual no ambiente Python são serializados para o disco, geralmente utilizando bibliotecas 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.

Fluxo de trabalho de registro baseado em serialização

A Databricks recomenda que você use o log baseado em código em vez do log baseado em serialização. Para obter detalhes sobre como utilizar o registro baseado em código, consulte Fluxo de trabalho de registro baseado em código.

Esta seção descreve como utilizar o log baseado em serialização.

Fluxo de trabalho de registro baseado em serialização com LangChain

Crie um notebook ou arquivo Python com seu código. O notebook ou arquivo deve conter uma cadeia LangChain, referenciada aqui como lc_chain.
Incluir mlflow.lang_chain.log_model(lc_model=lc_chain) no notebook ou arquivo.
Uma cópia serializada de PyFuncClass() é registrada em um modelo MLflow.
Implante o modelo.
Quando o ambiente de serviço é carregado, PyFuncClass é desserializado.
Quando chega uma solicitação de atendimento, lc_chain.invoke(...) é chamado.

Fluxo de trabalho de registro baseado em serialização com PyFunc

Crie um notebook ou arquivo Python com seu código. Para fins deste exemplo, o notebook ou arquivo é denominado notebook.py. O notebook ou arquivo deve conter uma classe PyFunc, referenciada aqui como PyFuncClass.
Inclua mlflow.pyfunc.log_model(python_model=PyFuncClass()) em notebook.py.
Uma cópia serializada de PyFuncClass() é registrada em um modelo MLflow.
Implante o modelo.
Quando o ambiente de serviço é carregado, PyFuncClass é desserializado.
Quando chega uma solicitação de atendimento, PyFuncClass.predict(...) é chamado.