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

MLflow As assinaturas de modelo definem os requisitos de esquema de entrada e saída para seu agente AI. A assinatura do modelo informa aos componentes internos e externos como interagir com seu agente. A assinatura do modelo atua como uma verificação de validação para garantir que as entradas atendam aos requisitos do esquema.

Por exemplo, para usar o aplicativo de avaliação do agente, seu agente deve seguir o esquema de entrada da avaliação do agente.

Esquemas de entrada suportados

O Mosaic AI Agent Framework é compatível com os seguintes esquemas de entrada.

Esquema de conclusão de bate-papo do OpenAI

Observação

A Databricks recomenda o esquema de conclusão de bate-papo da OpenAI porque ele é amplamente usado e é interoperável com muitas estruturas e aplicativos de agentes. Se o esquema de conclusão de bate-papo do OpenAI não atender às suas necessidades, você poderá definir seu próprio esquema. Consulte Esquemas de agentes personalizados.

• (Recomendado) A Databricks recomenda usar o esquema de conclusão de bate-papo do OpenAI. O esquema de conclusão de bate-papo do OpenAI deve ter uma matriz de objetos como parâmetro messages. Esse formato é melhor para aplicativos 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",
          },
      ]
  }
  

Solicitação de mensagem do Split Chat

SplitChatMessagesRequest é recomendado para aplicativos de bate-papo com vários turnos, especialmente quando o senhor deseja gerenciar a consulta atual e o histórico separadamente.

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

Linguagem de expressão Langchain

Se o seu agente usa LangChain, o senhor pode escrever sua cadeia em LangChain Expression Language. Em seu código de definição de cadeia, você pode usar um itemgetter para receber as mensagens ou os objetos query ou history, dependendo do formato de entrada.

Esquemas de saída suportados

O Mosaic AI Agent Framework é compatível com os seguintes esquemas de saída.

Resposta de conclusão do bate-papo

(Recomendado) O ChatCompletionResponse é recomendado para clientes com interoperabilidade de formato de resposta OpenAI.

LangChain - ChatCompletionsOutputParser

Se o seu agente usar LangChain, use ChatCompletionsOutputParser() de MLflow como sua cadeia de passos final. Isso formata a mensagem LangChain AI em um formato compatível com o agente.

  from mlflow.langchain.output_parsers import ChatCompletionsOutputParser

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

PyFunc - classes de entrada e saída adicionais

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.

  from mlflow.models.rag_signatures import ChatCompletionRequest, ChatCompletionResponse, ChainCompletionChoice, Message

  class RAGModel(PythonModel):
    ...
      def predict(self, context, model_input: ChatCompletionRequest) -> ChatCompletionResponse:
        ...
        return asdict(ChatCompletionResponse(
            choices=[ChainCompletionChoice(message=Message(content=text))]
        ))

Assinaturas explícitas e inferidas

O MLflow pode inferir o esquema de entrada e saída do seu agente em tempo de execução e criar uma assinatura automaticamente. Se você usa esquemas de entrada e saída compatíveis, as assinaturas inferidas são compatíveis com o Agent Framework. Para obter informações sobre os esquemas suportados, consulte Esquemas de entrada suportados.

No entanto, se você usar um esquema de agente personalizado, deverá definir explicitamente sua assinatura de modelo de acordo com as instruções em Esquemas de agente personalizados.

Esquemas de agentes personalizados

Você pode personalizar o esquema de um agente para passar e retornar campos adicionais de e para o agente criando uma subclasse de um esquema de entrada/saída compatível. Em seguida, adicione a chave extra custom_inputs e custom_outputs para conter os campos adicionais. Veja abaixo exemplos de código para Pyfunc e Langchain.

Para usar o databricks-agents SDK, as UIs de clientes do Databricks, como o AI Playground e o Review App, e outros recursos do Mosaic AI Agent Framework, o esquema do seu agente deve atender aos seguintes requisitos:

• O agente deve usar o site mlflow versão 2.17.1 ou acima.

• No notebook do agente, marque os campos adicionais adicionados em sua subclasse como Optional e atribua valores default.

• No Notebook do driver, construa um ModelSignature usando infer_signature com instâncias de suas subclasses.

• No Notebook do driver, crie um exemplo de entrada chamando asdict em sua subclasse.

Esquemas personalizados do PyFunc

Além dos requisitos acima, os agentes baseados em PyFunc também devem atender aos seguintes requisitos para interagir com o recurso de agente Mosaic AI.

Requisitos de esquema personalizado do PyFunc

No notebook do agente, as funções de previsão e previsão de transmissão devem atender aos seguintes requisitos:

• Tenha dicas de tipo para sua subclasse de entrada.

• Use a notação de ponto para acessar campos de classe de dados (por exemplo, use model_input.custom_input.id em vez de model_input["custom_inputs"]).

• Retorne um dictionary. Você pode chamar asdict em uma instância da sua subclasse para formatar o retorno como um dicionário.

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 Langchain

O Notebook a seguir mostra um exemplo de esquema personalizado usando LangChain. 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 Langchain Notebook

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

Driver de esquema personalizado Langchain Notebook

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