Criar e registrar agentes de IA

Artigo
06/21/2024

Importante

Esta funcionalidade está em Pré-visualização Pública.

Este artigo mostra como criar e registrar agentes de IA, como aplicativos RAG, usando 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, usar esses documentos para complementar o texto prompt e usar um modelo de base para resumir a resposta. O código que liga esses componentes, também chamados de etapas, é chamado de cadeia.

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

Com o Agent Framework, você pode usar quaisquer bibliotecas ou pacotes para criar código. O Agent Framework também facilita a iteração em seu código à medida que você o desenvolve e testa. Você pode configurar arquivos de configuração que permitem alterar parâmetros de código de forma rastreável sem ter que modificar o código real.

Requisitos

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

Esquema de entrada para o agente RAG

A seguir estão os formatos de entrada suportados para sua cadeia.

(Recomendado) Consultas usando o esquema de conclusão de chat do OpenAI. Ele deve ter uma matriz de objetos como parâmetro messages . Este formato é o 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 aplicativos de bate-papo com vários turnos, especialmente quando você 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 LangChain, Databricks recomenda escrever sua cadeia em LangChain Expression Language. Em seu código de definição de cadeia, você pode usar um itemgetter para obter as mensagens ou queryhistory objetos, dependendo do formato de entrada que você está usando.

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. Este 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 etapa final da cadeia. Sua saída deve retornar um único valor de cadeia de caracteres.

  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 usando o PyFunc, o Databricks recomenda o uso de dicas de tipo para anotar a predict() função 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 interna 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.

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

No Agent Framework, você pode usar parâmetros para controlar como os agentes são executados. Isso permite que você itere rapidamente variando as características do seu agente sem alterar o código. Os parâmetros são pares chave-valor que você define em um dicionário Python ou em um .yaml arquivo.

Para configurar o código, crie um ModelConfig, um conjunto de parâmetros chave-valor. O ModelConfig é um dicionário Python ou um .yaml arquivo. Por exemplo, você pode usar um dicionário durante o desenvolvimento e, em seguida, convertê-lo em um .yaml arquivo para implantação de produção e CI/CD. Para obter detalhes sobre ModelConfigo , 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')

Registrar o agente

O registro 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 que você possa 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 log baseado em serialização.

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

Fluxo de trabalho de log baseado em código

Para registro em log baseado em código, o código que registra seu agente ou cadeia deve estar em um bloco de anotações separado do código da cadeia. Este notebook é chamado de notebook de driver. Para obter um bloco de anotações de exemplo, consulte Blocos de anotações de exemplo.

Fluxo de trabalho de registro baseado em código com LangChain

Crie um bloco de anotações ou arquivo Python com seu código. Para fins deste exemplo, o bloco de anotações ou arquivo é chamado chain.py. O bloco de notas ou ficheiro deve conter uma cadeia LangChain, aqui referida como lc_chain.
Inclua mlflow.models.set_model(lc_chain) no bloco de anotações ou arquivo.
Crie um novo bloco de anotações para servir como o bloco de anotações de driver (chamado driver.py neste exemplo).
No caderno do controlador, inclua a chamada mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”). Esta chamada é executada chain.py e registra os resultados em um modelo MLflow.
Implemente o modelo.
Quando o ambiente de serviço é carregado, chain.py é executado.
Quando um pedido de serviço chega, lc_chain.invoke(...) é chamado.

Fluxo de trabalho de registro baseado em código com PyFunc

Crie um bloco de anotações ou arquivo Python com seu código. Para fins deste exemplo, o bloco de anotações ou arquivo é chamado chain.py. O bloco de anotações ou arquivo deve conter uma classe PyFunc, referida aqui como PyFuncClass.
Inclua mlflow.models.set_model(PyFuncClass) no bloco de anotações ou arquivo.
Crie um novo bloco de anotações para servir como o bloco de anotações de driver (chamado driver.py neste exemplo).
No caderno do controlador, inclua a chamada mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”). Esta chamada é executada chain.py e registra os resultados em um modelo MLflow.
Implemente o modelo.
Quando o ambiente de serviço é carregado, chain.py é executado.
Quando um pedido de serviço chega, PyFuncClass.predict(...) é chamado.

Código de exemplo para cadeias de registro

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 Catálogo Unity

Antes de implantar a cadeia, você deve registrá-la no Unity Catalog. Quando você registra a cadeia, ela é empacotada como um modelo no Unity Catalog, e você pode usar 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)

Exemplos de blocos de notas

Esses blocos de anotações criam uma cadeia simples "Olá, mundo" para ilustrar como criar um aplicativo em cadeia no Databricks. O primeiro exemplo cria uma cadeia simples. O segundo bloco de anotações de exemplo ilustra como usar parâmetros para minimizar alterações de código durante o desenvolvimento.

Registro em log baseado em código versus serialização

Para criar e registrar uma cadeia, você pode usar o log de MLflow baseado em código ou o log de MLflow baseado em serialização. O 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 para que ela possa ser invocada quando o ponto de extremidade é chamado.

Com o log MLflow baseado em serialização, o código da cadeia e o estado atual no ambiente Python são serializados em disco, geralmente usando 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 ponto de extremidade é 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 da 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 pode ser serializado.	`log_model(...)` deve ser chamado a partir de um notebook diferente do código da cadeia (chamado de notebook de driver).
Log de MLflow baseado em serialização	`log_model(...)` pode ser chamado a partir do mesmo notebook onde o modelo está definido.	* O código original não está disponível. * Todas as bibliotecas e objetos usados na cadeia devem suportar a serialização.

Fluxo de trabalho de log baseado em serialização

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

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

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

Crie um bloco de anotações ou arquivo Python com seu código. O bloco de notas ou ficheiro deve conter uma cadeia LangChain, aqui referida como lc_chain.
Inclua mlflow.lang_chain.log_model(lc_model=lc_chain) no bloco de anotações ou arquivo.
Uma cópia serializada do PyFuncClass() é registrada em um modelo MLflow.
Implemente o modelo.
Quando o ambiente de serviço é carregado, PyFuncClass é desserializado.
Quando um pedido de serviço chega, lc_chain.invoke(...) é chamado.

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

Crie um bloco de anotações ou arquivo Python com seu código. Para fins deste exemplo, o bloco de anotações ou arquivo é chamado notebook.py. O bloco de anotações ou arquivo deve conter uma classe PyFunc, referida aqui como PyFuncClass.
Incluir mlflow.pyfunc.log_model(python_model=PyFuncClass()) em notebook.py.
Uma cópia serializada do PyFuncClass() é registrada em um modelo MLflow.
Implemente o modelo.
Quando o ambiente de serviço é carregado, PyFuncClass é desserializado.
Quando um pedido de serviço chega, PyFuncClass.predict(...) é chamado.

Partilhar via