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

Esquema legado de entrada e saída do agente (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, configuração do servidor e fluxo de trabalho de implantação. Consulte Criar um agente de AI e implantá-lo no Databricks Apps. Para migrar um agente existente, consulte Migrar um agente do Model Serving para o Databricks Apps.

nota

A Databricks recomenda migrar para o esquema ResponsesAgent para criar agentes. Consulte Crie um agente AI e o implante no Databricks Apps.

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

Crie um agente ChatAgent legado​

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

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

Para saber como criar um(a) 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 esse 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 na Databricks. Em vez disso, basta envolver seu agente existente com a interface ChatAgent do MLflow:

  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 a implementação de 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 extras.

    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 o agente gerar uma resposta, converta a saída em um ou mais objetos ChatAgentMessage e os envolva em um ChatAgentResponse.

dica

Converter a saída do LangChain automaticamente

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

O que se segue é 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 seguintes notebooks mostram como criar ChatAgents de transmissão e não transmissão usando as bibliotecas populares OpenAI, LangGraph e AutoGen.

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

Agente de chamada de ferramentas do LangGraph

Abrir notebook em uma nova abaOpen in Databricks

Para saber como expandir os recursos desses agentes adicionando ferramentas, consulte Conecte agentes a ferramentas.

Respostas de transmissão do ChatAgent​

Agentes de transmissão entregam respostas em uma transmissão contínua de blocos menores e incrementais. A transmissão reduz a latência percebida e melhora a experiência do usuário para agentes conversacionais.

Para criar uma ChatAgent de transmissão, 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 ideal de transmissão ChatAgent nos documentos do MLflow.

O código a seguir mostra uma função predict_stream de exemplo; 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"]
      )

Crie um ChatModel agent legado​

importante

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

ChatModel é uma interface legada de autoria de agentes no MLflow que estende o esquema ChatCompletion do OpenAI, permitindo que você mantenha a compatibilidade com plataformas que suportam o padrão ChatCompletion enquanto adiciona funcionalidades personalizadas. Consulte MLflow: Para começar com o 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 requisição).

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

atenção

Os logs de solicitação e os 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ção sobre 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 a conclusão de 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. Notebooks fornecem um ambiente conveniente para desenvolver, testar e fazer iteração em seu agente.

A classe MyAgent estende mlflow.pyfunc.ChatModel, implementando o método predict necessá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 registra o agente no Model Registry e implanta o agente utilizando o Model Serving.

Esta separação segue o fluxo de trabalho recomendado pela Databricks para o registro de modelos usando a metodologia Modelos a partir de Código do MLflow.

Esquema de entrada SplitChatMessageRequest (obsoleto)​

SplitChatMessagesRequest permite que a query atual e a história sejam fornecidas 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 do StringResponse (obsoleto)​

StringResponse permite que você retorne a resposta do agente como um objeto com um único campo de string content:

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