Autor AI agentes em código

Este artigo mostra como criar um agente AI em código, usando MLflow ChatModel. Databricks aproveita MLflow ChatModel para garantir a compatibilidade com Databricks AI recurso de agente como avaliação, rastreamento e implantação.

O que é ChatModel?

ChatModel é uma classe MLflow projetada para simplificar a criação de agentes de conversação AI. Ele fornece uma interface padronizada para a criação de modelos compatíveis com a API ChatCompletion da OpenAI.

ChatModel estende o esquema ChatCompletion do OpenAI. Essa abordagem permite que você mantenha ampla compatibilidade com plataformas compatíveis com o padrão ChatCompletion, além de adicionar sua própria funcionalidade personalizada.

Ao usar ChatModel, os desenvolvedores podem criar agentes compatíveis com as ferramentas Databricks e MLflow para acompanhamento, avaliação e gerenciamento do ciclo de vida dos agentes, que são essenciais para a implantação de modelos prontos para produção.

MLflowConsulte: Começando a usar o ChatModel.

Requisitos

A Databricks recomenda instalar a versão mais recente do cliente MLflow Python ao desenvolver agentes.

Para criar e implantar agentes usando a abordagem deste artigo, o senhor deve atender aos seguintes requisitos:

  • Instale databricks-agents versão 0.15.0 e acima

  • Instalar mlflow versão 2.20.0 e acima

%pip install -U -qqqq databricks-agents>=0.15.0 mlflow>=2.20.0

Crie um agente do ChatModel

Você pode criar seu agente como uma subclasse de mlflow.pyfunc.ChatModel. Esse método oferece os seguintes benefícios:

  • Permite que o senhor escreva código de agente compatível com o esquema ChatCompletion usando classes Python tipadas.

  • O MLflow inferirá automaticamente uma assinatura compatível com a conclusão do chat ao registrar o agente, mesmo sem um input_example. Isso simplifica o processo de registro e implantação do agente. Consulte Inferir assinatura do modelo durante o registro.

O código a seguir é melhor executado em um notebook Databricks. O Notebook oferece um ambiente conveniente para o desenvolvimento, teste e iteração de seu agente.

A classe MyAgent estende mlflow.pyfunc.ChatModel, implementando o método predict necessário. Isso garante a compatibilidade com o Mosaic AI Agent Framework.

A classe também inclui os métodos opcionais _create_chat_completion_chunk e predict_stream para lidar com as saídas de transmissão.

from dataclasses import dataclass
from typing import Optional, Dict, List, Generator
from mlflow.pyfunc import ChatModel
from mlflow.types.llm import (
    # Non-streaming helper classes
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionChunk,
    ChatMessage,
    ChatChoice,
    ChatParams,
    # Helper classes for streaming agent output
    ChatChoiceDelta,
    ChatChunkChoice,
)

class MyAgent(ChatModel):
    """
    Defines a custom agent that processes ChatCompletionRequests
    and returns ChatCompletionResponses.
    """
    def predict(self, context, messages: list[ChatMessage], params: ChatParams) -> ChatCompletionResponse:
        last_user_question_text = messages[-1].content
        response_message = ChatMessage(
            role="assistant",
            content=(
                f"I will always echo back your last question. Your last question was: {last_user_question_text}. "
            )
        )
        return ChatCompletionResponse(
            choices=[ChatChoice(message=response_message)]
        )

    def _create_chat_completion_chunk(self, content) -> ChatCompletionChunk:
        """Helper for constructing a ChatCompletionChunk instance for wrapping streaming agent output"""
        return ChatCompletionChunk(
                choices=[ChatChunkChoice(
                    delta=ChatChoiceDelta(
                        role="assistant",
                        content=content
                    )
                )]
            )

    def predict_stream(
        self, context, messages: List[ChatMessage], params: ChatParams
    ) -> Generator[ChatCompletionChunk, None, None]:
        last_user_question_text = messages[-1].content
        yield self._create_chat_completion_chunk(f"Echoing back your last question, word by word.")
        for word in last_user_question_text.split(" "):
            yield self._create_chat_completion_chunk(word)

agent = MyAgent()
model_input = ChatCompletionRequest(
    messages=[ChatMessage(role="user", content="What is Databricks?")]
)
response = agent.predict(context=None, model_input=model_input)
print(response)

Embora a classe de agente MyAgent seja definida em um Notebook, o senhor deve criar um Notebook de driver separado. O driver Notebook logs o agente para Model Registry e implantado o agente usando o modelo servindo.

Essa separação segue o fluxo de trabalho recomendado pela Databricks para modelos de registro usando a metodologia Models from Code do MLflow.

Exemplo: Envolver LangChain em ChatModel

Se o senhor tiver um modelo LangChain existente e quiser integrá-lo a outro recurso de agente Mosaic AI, poderá envolvê-lo em um MLflow ChatModel para garantir a compatibilidade.

Este exemplo de código executa as seguintes etapas para envolver um runnable LangChain como ChatModel:

  1. Envolva a saída final da LangChain com mlflow.langchain.output_parsers.ChatCompletionOutputParser para produzir uma assinatura de saída de conclusão de bate-papo

  2. A classe LangchainAgent estende mlflow.pyfunc.ChatModel e implementa dois métodos key:

    • predict: processa previsões síncronas invocando a cadeia e retornando uma resposta formatada.

    • predict_stream: Trata as previsões de transmissão invocando a cadeia e produzindo pedaços de respostas.

from mlflow.langchain.output_parsers import ChatCompletionOutputParser
from mlflow.pyfunc import ChatModel
from typing import Optional, Dict, List, Generator
from mlflow.types.llm import (
    ChatCompletionResponse,
    ChatCompletionChunk
)

chain = (
    <your chain here>
    | ChatCompletionOutputParser()
)

class LangchainAgent(ChatModel):
    def _prepare_messages(self, messages: List[ChatMessage]):
        return {"messages": [m.to_dict() for m in messages]}

    def predict(
        self, context, messages: List[ChatMessage], params: ChatParams
    ) -> ChatCompletionResponse:
        question = self._prepare_messages(messages)
        response_message = self.chain.invoke(question)
        return ChatCompletionResponse.from_dict(response_message)

    def predict_stream(
        self, context, messages: List[ChatMessage], params: ChatParams
    ) -> Generator[ChatCompletionChunk, None, None]:
        question = self._prepare_messages(messages)
        for chunk in chain.stream(question):
          yield ChatCompletionChunk.from_dict(chunk)

Use parâmetros para configurar o agente

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, crie um ModelConfig, um conjunto de parâmetros key-value. ModelConfig é um dicionário Python ou um arquivo .yaml. Por exemplo, o senhor pode usar 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')

Definir esquema do recuperador

AI Os agentes costumam usar recuperadores, um tipo de ferramenta de agente que localiza e retorna documentos relevantes usando um índice de pesquisa vetorial. Para obter mais informações sobre recuperadores, consulte Recuperação não estruturada AI agent tools.

Para garantir que os recuperadores sejam rastreados corretamente, chame mlflow.models.set_retriever_schema quando você define seu agente em código. Use set_retriever_schema para mapear os nomes das colunas na tabela retornada para os campos esperados do MLflow, como primary_key, text_column e doc_uri.

# Define the retriever's schema by providing your column names
# These strings should be read from a config dictionary
mlflow.models.set_retriever_schema(
    name="vector_search",
    primary_key="chunk_id",
    text_column="text_column",
    doc_uri="doc_uri"
    # other_columns=["column1", "column2"],
)

Observação

A coluna doc_uri é especialmente importante ao avaliar o desempenho do retriever. doc_uri é o identificador principal dos documentos retornados pelo recuperador, permitindo compará-los com os conjuntos de avaliação da verdade básica. Consulte Conjuntos de avaliação.

Você também pode especificar colunas adicionais no esquema do seu retriever fornecendo uma lista de nomes de colunas com o campo other_columns.

Se você tiver vários recuperadores, poderá definir vários esquemas usando nomes exclusivos para cada esquema de recuperação.

Entradas e saídas personalizadas

Alguns cenários podem exigir entradas adicionais dos agentes, como client_type e session_id, ou saídas como links de fontes de recuperação que não devem ser incluídos no histórico do chat para interações futuras.

Para esses cenários, o MLflow ChatModel suporta nativamente o aumento das solicitações e respostas de conclusão de bate-papo da OpenAI com os campos ChatParams custom_input e custom_output.

Veja os exemplos a seguir para saber como criar entradas e saídas personalizadas para agentes PyFunc e LangGraph.

Aviso

Atualmente, o aplicativo Agent Evaluation Review não oferece suporte à renderização de rastreamentos para agentes com campos de entrada adicionais.

Esquemas personalizados do PyFunc

O Notebook a seguir mostra um exemplo de esquema personalizado usando PyFunc.

Agente de esquema personalizado PyFunc Notebook

Abra o bloco de anotações em outra guia

Notebook do driver de esquema personalizado PyFunc

Abra o bloco de anotações em outra guia

Esquemas personalizados do LangGraph

O Notebook a seguir mostra um exemplo de esquema personalizado usando o LangGraph. O senhor pode modificar a função wrap_output no Notebook para analisar e extrair informações da transmissão da mensagem.

Agente de esquema personalizado LangGraph Notebook

Abra o bloco de anotações em outra guia

Driver de esquema personalizado LangGraph Notebook

Abra o bloco de anotações em outra guia

Fornecer custom_inputs no AI Playground e no aplicativo de revisão de agentes

Se o seu agente aceitar entradas adicionais usando o campo custom_inputs, o senhor poderá fornecer manualmente essas entradas no AI Playground e no aplicativo de avaliação do agente.

  1. No AI Playground ou no aplicativo Agent Review, selecione o ícone de engrenagem Ícone de engrenagem.

  2. Habilite custom_inputs.

  3. Forneça um objeto JSON que corresponda ao esquema de entrada definido pelo seu agente.

    Forneça custom_inputs no playground AI.

transmissão propagação de erros

Mosaic AI propagação de quaisquer erros encontrados durante a transmissão com os últimos tokens em databricks_output.error. Cabe ao cliente chamador tratar e revelar adequadamente esse erro.

{
  "delta": …,
  "databricks_output": {
    "trace": {...},
    "error": {
      "error_code": BAD_REQUEST,
      "message": "TimeoutException: Tool XYZ failed to execute"
    }
  }
}

Notebooks de Exemplo

Esses Notebooks criam uma cadeia simples de "Hello, world" para ilustrar a criação de um agente em Databricks. O primeiro exemplo cria uma cadeia simples, e o segundo exemplo do Notebook ilustra como usar parâmetros para minimizar as alterações de código durante o desenvolvimento.

Caderno de cadeia simples

Abra o bloco de anotações em outra guia

Notebook simples com acionador de cadeia

Abra o bloco de anotações em outra tab

Caderno de cadeia parametrizado

Abra o bloco de anotações em outra tab

Notebook do driver de cadeia parametrizado

Abra o bloco de anotações em outra tab