Pular para o conteúdo principal
Última atualização em

Esquema de agente de entrada e saída legado (Model Serving)

info

Para novos casos de uso, a Databricks recomenda implantar agentes no Databricks Apps para controle total sobre o código do agente, a configuração do servidor e o fluxo de trabalho de implantação. Consulte Crie um agente de AI e o tenha implantado no Databricks Apps. Para migrar um agente existente, consulte Migre um agente do Model Serving para o Databricks Apps.

nota

A Databricks recomenda migrar para o esquema ResponsesAgent para criar agentes. Consulte Criar um agente de AI e implantá-lo no Databricks Apps.

Agentes de AI devem aderir a requisitos específicos de esquema de entrada e saída para serem compatíveis com outros recursos no Databricks. Esta página explica como usar as assinaturas e interfaces legadas de autoria de agente: interface ChatAgent, interface ChatModel, o esquema de entrada SplitChatMessageRequest e o esquema de saída StringResponse.

Criar um agente ChatAgent legado​

A interface ChatAgent do MLflow é similar a, mas não estritamente compatível com, o esquema ChatCompletion da OpenAI.

O ChatAgent integra facilmente agentes existentes para compatibilidade com o Databricks.

Para saber como criar um ChatAgent, consulte os exemplos na seção a seguir e a documentação do MLflow - O que é a interface ChatAgent.

Para criar e implantar agentes usando ChatAgent, instale o seguinte:

  • databricks-agents0.16.0 ou acima
  • mlflow 2.20.2 ou acima
  • Python 3.10 ou acima.
    • Para atender a este requisito, você pode usar compute serverless ou Databricks Runtime 13.3 LTS ou acima.
Python
%pip install -U -qqqq databricks-agents==0.16.0 mlflow==2.20.2

E se eu já tiver um agente?​

Se você já tem um agente construído com LangChain, LangGraph ou um framework similar, você não precisa reescrever seu agente para usá-lo no Databricks. Em vez disso, basta envolver seu agente existente com a interface MLflow ChatAgent:

  1. Escreva uma classe wrapper Python que herda de mlflow.pyfunc.ChatAgent.

    Dentro da classe wrapper, mantenha seu agente existente como um atributo self.agent = your_existing_agent.

  2. A classe ChatAgent exige que você implemente um método predict para lidar com solicitações sem transmissão.

    predict deve aceitar:

    • messages: list[ChatAgentMessage], que é uma lista de ChatAgentMessage cada um com uma função (como "usuário" ou "assistente"), o prompt e um ID.

    • (Opcional) context: Optional[ChatContext] e custom_inputs: Optional[dict] para dados adicionais.

    Python
    import uuid

    # input example
    [
      ChatAgentMessage(
        id=str(uuid.uuid4()),  # Generate a unique ID for each message
        role="user",
        content="What's the weather in Paris?"
      )
    ]

    predict deve retornar um ChatAgentResponse.

    Python
    import uuid

    # output example
    ChatAgentResponse(
      messages=[
        ChatAgentMessage(
          id=str(uuid.uuid4()),  # Generate a unique ID for each message
          role="assistant",
          content="It's sunny in Paris."
        )
      ]
    )

  3. Converter entre formatos

    Em predict, converta as mensagens recebidas de list[ChatAgentMessage] para o formato de entrada que seu agente espera.

    Depois que seu agente gera uma resposta, converta sua saída para um ou mais objetos ChatAgentMessage e envolva-os em um(a) ChatAgentResponse.

dica

Converta a saída do LangChain automaticamente

Se você estiver encapsulando um agente LangChain, poderá usar mlflow.langchain.output_parsers.ChatAgentOutputParser para converter automaticamente as saídas do LangChain no esquema MLflow ChatAgentMessage e ChatAgentResponse.

O seguinte é um padrão simplificado para converter seu agente:

Python
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatAgentChunk
import uuid


class MyWrappedAgent(ChatAgent):
  def __init__(self, agent):
    self.agent = agent

  def predict(self, messages, context=None, custom_inputs=None):
    # Convert messages to your agent's format
    agent_input = ... # build from messages
    agent_output = self.agent.invoke(agent_input)
    # Convert output to ChatAgentMessage
    return ChatAgentResponse(
      messages=[ChatAgentMessage(role="assistant", content=agent_output, id=str(uuid.uuid4()),)]
    )

  def predict_stream(self, messages, context=None, custom_inputs=None):
    # If your agent supports streaming
    for chunk in self.agent.stream(...):
      yield ChatAgentChunk(delta=ChatAgentMessage(role="assistant", content=chunk, id=str(uuid.uuid4())))

Para exemplos completos, consulte os Notebooks na seção a seguir.

ChatAgent exemplos​

Os notebooks a seguir mostram como criar ChatAgents com transmissão e sem transmissão usando as bibliotecas populares OpenAI, LangGraph e AutoGen.

Se você estiver encapsulando um agente LangChain, poderá usar mlflow.langchain.output_parsers.ChatAgentOutputParser para converter automaticamente as saídas do LangChain no esquema MLflow ChatAgentMessage e ChatAgentResponse.

Agente de chamada de ferramentas LangGraph

Abrir notebook em uma nova aba

Para aprender como expandir as capacidades desses agentes adicionando ferramentas, veja Conectar agentes a ferramentas.

Respostas do ChatAgent em transmissão​

Agentes de transmissão fornecem respostas em um fluxo contínuo de partes incrementais menores. A transmissão reduz a latência percebida e melhora a experiência do usuário para agentes conversacionais.

Para criar uma transmissão ChatAgent, defina um método predict_stream que retorna um gerador que produz objetos ChatAgentChunk — cada ChatAgentChunk contém uma parte da resposta. Leia mais sobre o comportamento de transmissão ChatAgent ideal na documentação do MLflow.

O código a seguir mostra uma função de exemplo predict_stream; para exemplos completos de agentes de transmissão, consulte exemplos de ChatAgent:

Python
def predict_stream(
  self,
  messages: list[ChatAgentMessage],
  context: Optional[ChatContext] = None,
  custom_inputs: Optional[dict[str, Any]] = None,
) -> Generator[ChatAgentChunk, None, None]:
  # Convert messages to a format suitable for your agent
  request = {"messages": self._convert_messages_to_dict(messages)}

  # Stream the response from your agent
  for event in self.agent.stream(request, stream_mode="updates"):
    for node_data in event.values():
      # Yield each chunk of the response
      yield from (
        ChatAgentChunk(**{"delta": msg}) for msg in node_data["messages"]
      )

Autorar um agente de ChatModel legado​

importante

A Databricks recomenda a interface ChatAgent para criar agentes ou aplicativos de AI generativa. Para migrar do ChatModel para o ChatAgent, consulte a documentação do MLflow - Migrar do ChatModel para o ChatAgent.

ChatModel é uma interface de autoria de agente legada no MLflow que estende o esquema de ChatCompletion do OpenAI, permitindo manter a compatibilidade com plataformas que suportam o padrão ChatCompletion enquanto adiciona funcionalidade personalizada. Consulte MLflow: Introdução ao ChatModel para obter detalhes adicionais.

A criação do seu agente como uma subclasse de mlflow.pyfunc.ChatModel oferece os seguintes benefícios:

  • Habilita a saída de transmissão do agente ao invocar um agente servido (ignorando {stream: true} no corpo da solicitação).

  • Habilita automaticamente as tabelas de inferência do AI Gateway quando seu agente é servido, fornecendo acesso a metadados aprimorados do log de solicitação, como o nome do solicitante.

atenção

Logs de solicitação e logs de avaliação foram descontinuados e serão removidos em uma versão futura. Consulte a descontinuação dos logs de solicitação e de avaliação para obter orientações de migração.

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

  • O MLflow infere automaticamente uma assinatura compatível com preenchimento de chat ao registrar o agente, mesmo sem um(a) 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. Notebooks fornecem um ambiente conveniente para desenvolver, testar e iterar em seu agente.

A classe MyAgent estende mlflow.pyfunc.ChatModel, implementando o método predict obrigatório. Isso garante compatibilidade com Agentes Personalizados.

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

Python
# Install a pinned version of mlflow
%pip install -U mlflow==2.20.2
dbutils.library.restartPython()
Python
import re
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 re.findall(r"\S+\s*", last_user_question_text):
      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, messages=model_input.messages, params=None)
print(response)

Embora você defina a classe de agente MyAgent em um notebook, recomendamos criar um notebook de driver separado. O Notebook do driver faz o log do agente para o Model Registry e implanta o agente usando o Model Serving.

Essa separação segue o fluxo de trabalho recomendado pela Databricks para registrar modelos usando a metodologia Modelos por Código do MLflow.

Esquema de entrada SplitChatMessageRequest (obsoleto)​

SplitChatMessagesRequest permite passar a query atual e a história separadamente como entrada do agente.

Python
  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 retornar a resposta do agente como um objeto com um único campo content de strings:

{"content": "This is an example string response"}
Nesta página