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

Os agentes de IA devem aderir a requisitos específicos de entrada e saída schema para serem compatíveis com outros recursos no Azure Databricks. Este artigo explica como garantir que o seu agente de IA cumpra esses requisitos e como personalizar os dados de entrada e saída do agente schema, garantindo a compatibilidade.

O Mosaic AI usa de assinaturas de modelo MLflow para definir os requisitos de schema 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 que eles aderem ao schema.

conclusão do chat OpenAI schema (recomendado)

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

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

O MLflow fornece APIs convenientes para agentes do tipo LangChain e PyFunc, ajudando-o a criar agentes compatíveis com a conclusão do bate-papo schema.

Implemente a conclusão do bate-papo com LangChain

Se o seu agente usa LangChain, utilize o ChatCompletionOutputParser() do MLflow para formatar a saída final do seu agente para garantir compatibilidade com a finalização da conversa schema. Se você usa o LangGraph, consulte 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 você usa o PyFunc, o Databricks recomenda escrever seu agente como uma subclasse de mlflow.pyfunc.ChatModel. Este método oferece os seguintes benefícios:

• Permite escrever código de agente compatível com o schema para conclusão de conversas usando classes Python tipadas.

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

  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 às conclusões de bate-papo do OpenAI schema para a maioria dos casos de utilização de agentes.

No entanto, alguns cenários podem exigir entradas adicionais, como client_type e session_id, ou saídas, como links de origem de recuperação que não devem ser incluídos no histórico de bate-papo para interações futuras.

Para esses cenários, o Mosaic AI Agent Framework suporta ampliar os pedidos de conclusão de chat e as respostas 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.

Esquemas personalizados do PyFunc

Os blocos de anotações a seguir mostram um exemplo de schema personalizado usando o PyFunc.

Caderno personalizado do agente PyFunc schema

Get portátil

Notebook do driver PyFunc personalizado schema

Get portátil

esquemas personalizados LangGraph

Os cadernos a seguir mostram um exemplo personalizado de schema usando o LangGraph. Você pode modificar a função wrap_output nos blocos de anotações para analisar e extrair informações do fluxo de mensagens.

Notebook de agente personalizado schema do LangGraph

Get portátil

LangGraph notebook driver de schema personalizado

Get caderno

Forneça custom_inputs no AI Playground e no aplicativo de revisão de agentes

Se o agente aceitar entradas adicionais usando o campo custom_inputs, você poderá inserir manualmente estas entradas no AI Playground e no aplicativo de revisão do agente .

1. Tanto no AI Playground como na aplicação Agent Review, select o ícone de engrenagem .

2. Habilite custom_inputs.

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

   

Esquemas de entrada e saída herdados

Os parâmetros de entrada SplitChatMessageRequest schema e de saída StringResponse schema atingiram o fim de vida útil. Se você estiver usando qualquer um desses esquemas herdados, o Databricks recomenda que você migre para o schemade conclusão de bate-papo recomendado.

SplitChatMessageRequest de entrada schema (obsoleto)

SplitChatMessagesRequest permite-te passar a consulta atual e o histórico separadamente como entrada para o agente.

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

Saída de StringResponse schema (preterido)

StringResponse permite-lhe devolver a resposta do agente como um objeto com um único campo de cadeia de caracteres content.

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