Definir o esquema de entrada e saída de um agente

AI Os agentes devem aderir a requisitos específicos de esquema de entrada e saída para serem compatíveis com outros recursos em Databricks. Este artigo explica como garantir que o agente AI atenda a esses requisitos e como personalizar o esquema de entrada e saída do agente, garantindo a compatibilidade.

O Mosaic AI usa MLflow Model Signatures para definir os requisitos do esquema de entrada e saída de um agente. A assinatura do modelo informa aos componentes internos e externos como interagir com seu agente e valida se eles aderem ao esquema.

Esquema de conclusão de bate-papo do OpenAI (recomendado)

A Databricks recomenda usar o esquema de conclusão de bate-papo da OpenAI para definir a entrada e a saída do agente. Esse esquema é amplamente adotado e compatível com muitas estruturas e aplicativos de agentes, incluindo os da Databricks.

Consulte a documentação do OpenAI para obter mais informações sobre o esquema de entrada e o esquema de saída do preenchimento de bate-papo.

Observação

O esquema de conclusão de bate-papo do OpenAI é simplesmente um padrão para estruturar as entradas e saídas do agente. A implementação desse esquema não envolve o uso de recursos ou modelos da OpenAI.

O MLflow oferece APIs convenientes para agentes LangChain e PyFunc, ajudando o senhor a criar agentes compatíveis com o esquema de conclusão de bate-papo.

Implementar o preenchimento de bate-papo com LangChain

Se o seu agente usar LangChain, use o site ChatCompletionOutputParser() do MLflow para formatar a saída final do agente para que seja compatível com o esquema de conclusão do chat. Se você usa o LangGraph, consulte os esquemas personalizados do LangGraph.

  from mlflow.langchain.output_parsers import ChatCompletionOutputParser

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

Implemente a conclusão do bate-papo com o PyFunc

Se o senhor usar o PyFunc, a Databricks recomenda escrever 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 de conclusão de chat usando classes Python tipadas.

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

  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)
  

Entradas e saídas personalizadas

A Databricks recomenda aderir ao esquema de conclusões de bate-papo da OpenAI para a maioria dos casos de uso de agentes.

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

Para esses cenários, o Mosaic AI Agent Framework suporta o aumento das solicitações e respostas de conclusão de bate-papo do OpenAI com entradas e saídas personalizadas.

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 Copiar link para importação

Notebook do driver de esquema personalizado PyFunc

Abra o bloco de anotações em outra guia Copiar link para importação

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 Copiar link para importação

Driver de esquema personalizado LangGraph Notebook

Abra o bloco de anotações em outra guia Copiar link para importação

Forneça custom_inputs no AI Playground e no aplicativo de avaliaçã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 .

2. Habilite custom_inputs.

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

   

Esquemas de entrada e saída legados

O esquema de entrada SplitChatMessageRequest e o esquema de saída StringResponse foram descontinuados. Se o senhor estiver usando um desses esquemas antigos, a Databricks recomenda que migre para o esquema de conclusão de bate-papo recomendado.

Esquema de entrada SplitChatMessageRequest (obsoleto)

SplitChatMessagesRequest permite que o senhor passe a consulta atual e o histórico separadamente como entrada do agente.

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

Esquema de saída StringResponse (obsoleto)

StringResponse permite que o senhor retorne a resposta do agente como um objeto com um único campo de strings content:

{"content": "This is an example string response"}