Esquema de agente de entrada e saída herdado

Os agentes de IA devem seguir 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 de criação do agente herdado: ChatAgent interface, ChatModel interface, esquema SplitChatMessageRequest de entrada e esquema StringResponse de saída.

Criar um agente herdado do ChatAgent

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

Para saber como criar um ChatAgent, veja 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-agents 0.16.0 ou superior
• mlflow 2.20.2 ou superior
• Python 3.10 ou superior.
  • Para atender a esse requisito, você pode usar computação sem servidor ou Databricks Runtime 13.3 LTS ou superior.

%pip install -U -qqqq databricks-agents>=0.16.0 mlflow>=2.20.2

E se eu já tiver um agente?

Se você já tiver um agente criado com LangChain, LangGraph ou uma estrutura semelhante, você não precisará reescrever seu agente para usá-lo no Databricks. Em vez disso, basta encapsular seu agente existente com a interface MLflow ChatAgent :

1. Escreva uma classe de wrapper do 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 ChatAgent classe exige que você implemente um predict método para lidar com solicitações que não são de streaming.

   predict deve aceitar:

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

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

   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.

   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 que seu agente espera.

   Depois que o agente gerar uma resposta, converta a sua saída em um ou mais objetos ChatAgentMessage e envolva-os em um ChatAgentResponse.

Veja a seguir um modelo simplificado para converter seu agente:

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 obter exemplos completos, consulte os blocos de anotações na seção a seguir.

ChatAgent Exemplos

Os blocos de anotações a seguir mostram como criar streaming e não streaming ChatAgents usando as bibliotecas populares OpenAI, LangGraph e AutoGen.

LangGraph

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 do LangGraph

Obter laptop

OpenAI

Agente de chamada de ferramentas do OpenAI

Obter laptop

Agente de chamadas de ferramentas da API de Respostas do OpenAI

Obter laptop

Agente de chat exclusivo do OpenAI

Obter laptop

AutoGen

Agente de ativação de ferramentas do AutoGen

Obter laptop

DSPy

Agente somente chat do DSPy

Obter laptop

Para saber como expandir os recursos desses agentes adicionando ferramentas, consulte ferramentas do agente de IA.

Respostas de ChatAgent de streaming

Os agentes de streaming fornecem respostas em um fluxo contínuo de partes menores e incrementais. O streaming reduz a latência percebida e melhora a experiência do usuário para agentes de conversação.

Para criar um streaming ChatAgent, defina um método predict_stream que retorne um gerador que produza objetos ChatAgentChunk – cada um ChatAgentChunk contendo uma parte da resposta. Leia mais sobre o comportamento ideal ChatAgent de streaming nos documentos do MLflow.

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

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"]
      )

Criar um agente de ChatModel herdado

ChatModel é uma interface de criação de agente herdada no MLflow que estende o esquema ChatCompletion do OpenAI, permitindo que você mantenha a compatibilidade com plataformas que dão suporte ao padrão ChatCompletion ao adicionar funcionalidade personalizada. Consulte MLflow: Introdução ao ChatModel para obter detalhes adicionais.

Criar seu agente como uma subclasse de mlflow.pyfunc.ChatModel oferece os seguintes benefícios:

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

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

  Aviso

  Os logs de solicitação e os logs de avaliação foram preteridos e serão removidos em uma versão futura. Consulte os logs de solicitação e a substituição de logs de avaliação para obter diretrizes de migração.

• Permite que você escreva código de agente compatível com o esquema ChatCompletion usando classes python tipadas.

• O MLflow deduz automaticamente uma assinatura compatível com a conclusão do chat, mesmo sem um input_example, ao registrar o agente. Isso simplifica o processo de registro e implantação do agente. Confira Inferir assinatura de modelo durante o registro em log.

O código a seguir é melhor executado em um notebook do Databricks. Os notebooks fornecem um ambiente conveniente para desenvolver, testar e iterar em seu agente.

A classe MyAgent estende mlflow.pyfunc.ChatModel, implementando o método de predict necessário. Isso garante a compatibilidade com o Mosaic AI Agent Framework.

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

# Install the latest version of mlflow
%pip install -U mlflow
dbutils.library.restartPython()

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 MyAgent de agente em um bloco de anotações, recomendamos a criação de um bloco de anotações de driver separado. O notebook de controle registra o agente no Model Registry e implanta o agente usando o Model Serving.

Essa separação segue o fluxo de trabalho recomendado pelo Databricks para registrar modelos em log usando a metodologia 'Models from Code' do MLflow.

Esquema de entrada SplitChatMessageRequest (preterido)

SplitChatMessagesRequest permite que você passe a consulta e o histórico atuais separadamente para o agente como entrada.

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

Esquema de saída StringResponse (preterido)

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

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