Esquema heredado de agentes de entrada y salida (Modelo de servicio)

Los agentes deben cumplir los requisitos específicos del esquema de entrada y salida para que sean compatibles con otras características de Databricks. En esta página se explica cómo utilizar las firmas e interfaces heredadas de creación de agentes: la ChatAgent interfaz, la ChatModel interfaz, el SplitChatMessageRequest esquema de entrada y el StringResponse esquema de salida.

Crear un agente ChatAgent heredado

La interfaz MLflow ChatAgent es similar a, pero no es estrictamente compatible con, el esquema openAI ChatCompletion .

Para obtener información sobre cómo crear un ChatAgent, consulte los ejemplos en la sección siguiente y la documentación de MLflow - ¿Qué es la interfaz de ChatAgent?

Para crear e implementar agentes mediante ChatAgent, instale lo siguiente:

• databricks-agents 0.16.0 o superior
• mlflow 2.20.2 o superior
• Python 3.10 o superior.
  • Para cumplir este requisito, puede utilizar informática sin servidor o Databricks Runtime 13.3 LTS o superior.

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

¿Qué ocurre si ya tengo un agente?

Si ya tiene un agente compilado con LangChain, LangGraph o un marco similar, no es necesario volver a escribir el agente para usarlo en Databricks. En su lugar, encapsule el agente existente con la interfaz MLflow ChatAgent :

1. Escriba una clase contenedora de Python que herede de mlflow.pyfunc.ChatAgent.

   Dentro de la clase contenedora, mantenga el agente existente como atributo self.agent = your_existing_agent.

2. La ChatAgent clase requiere que implemente un predict método para controlar las solicitudes que no son de streaming.

   predict debe aceptar:

   • messages: list[ChatAgentMessage], que es una lista de ChatAgentMessage, cada uno de los elementos con un rol (como "usuario" o "asistente"), la instrucción y un ID.

   • (Opcional) context: Optional[ChatContext] y custom_inputs: Optional[dict] para datos adicionales.

   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 debe devolver un 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. Conversión entre formatos

   En predict, convierta los mensajes entrantes de list[ChatAgentMessage] al formato de entrada que espera su agente.

   Después de que el agente genere una respuesta, conviértalo en uno o varios objetos ChatAgentMessage y encapsúlelos en un ChatAgentResponse.

Consulte la siguiente plantilla simplificada para convertir el 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 obtener ejemplos completos, consulte los cuadernos de la sección siguiente.

ejemplos de ChatAgent

En los cuadernos siguientes se muestra cómo crear ChatAgents con flujos o sin flujos mediante las bibliotecas populares de OpenAI, LangGraph y AutoGen.

LangGraph

Si estás encapsulando un agente de LangChain, puedes usar mlflow.langchain.output_parsers.ChatAgentOutputParser para convertir automáticamente los resultados de LangChain en el esquema de MLflow ChatAgentMessage y ChatAgentResponse.

Agente de invocación de herramientas de LangGraph

Obtener el portátil

OpenAI

Agente de llamada a herramientas de OpenAI

Obtener el portátil

Agente de llamada a herramientas de La API de respuestas de OpenAI

Obtener el portátil

Agente de solo chat de OpenAI

Obtener el portátil

AutoGen

Agente de invocación de herramientas de AutoGen

Obtener el portátil

DSPy

Agente de solo chat de DSPy

Obtener el portátil

Para aprender cómo ampliar las capacidades de estos agentes añadiendo herramientas, consulta MCPs y herramientas para agentes.

Respuestas en streaming de ChatAgent

Los agentes de streaming proporcionan respuestas en un flujo continuo de fragmentos incrementales más pequeños. El streaming reduce la latencia percibida y mejora la experiencia del usuario para los agentes conversacionales.

Para crear un streaming ChatAgent, defina un predict_stream método que devuelva un generador que produce ChatAgentChunk objetos, cada uno ChatAgentChunk contiene una parte de la respuesta. Obtenga más información sobre el comportamiento de streaming ideal ChatAgent en los documentos de MLflow.

En el código siguiente se muestra una función de ejemplo predict_stream , para ver ejemplos completos de agentes de streaming, vea Ejemplos 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"]
      )

Crear un agente ChatModel heredado

ChatModel es una interfaz de creación de agentes heredado en MLflow que amplía el esquema ChatCompletion de OpenAI, lo que le permite mantener la compatibilidad con las plataformas que admiten el estándar ChatCompletion al agregar funcionalidad personalizada. Consulte MLflow: Getting Started with ChatModel (Introducción a ChatModel ) para obtener más información.

La creación del agente como una subclase de mlflow.pyfunc.ChatModel proporciona las siguientes ventajas:

• Permite la transmisión de la salida del agente al invocar un agente hospedado (omitiendo {stream: true} en el cuerpo de la solicitud).

• Habilita automáticamente las tablas de inferencia de AI Gateway cuando se atiende el agente, lo que proporciona acceso a metadatos de registro de solicitudes mejorados, como el nombre del solicitante.

  Warning

  Los registros de solicitud y los registros de evaluación están en desuso y se quitarán en una versión futura. Consulta la obsolescencia de los registros de solicitudes y de evaluación para obtener orientación sobre la migración.

• Permite escribir código de agente compatible con el esquema de ChatCompletion mediante clases tipadas de Python.

• MLflow deduce automáticamente una firma compatible con la finalización del chat al registrar el agente, incluso sin un input_example. Esto simplifica el proceso de registro e implementación del agente. Consulte Deducción de la firma del modelo durante el registro.

El código siguiente se ejecuta mejor en un cuaderno de Databricks. Los cuadernos ofrecen un entorno práctico para desarrollar, probar y perfeccionar tu agente.

La clase MyAgent extiende mlflow.pyfunc.ChatModel, implementando el método predict necesario. Esto garantiza la compatibilidad con los agentes personalizados.

La clase también incluye los métodos opcionales _create_chat_completion_chunk y predict_stream para controlar las salidas de streaming.

# Install a pinned version of mlflow
%pip install -U mlflow==2.20.2
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)

Aunque defina la clase MyAgent del agente en un cuaderno, se recomienda crear un cuaderno de controladores independiente. El cuaderno del controlador registra el agente en el Registro de modelos e implementa el agente mediante Model Serving.

Esta separación sigue el flujo de trabajo recomendado por Databricks para registrar modelos usando la metodología Models from Code de MLflow.

Esquema de entrada SplitChatMessageRequest (en desuso)

SplitChatMessagesRequest le permite pasar la consulta actual y el historial por separado como entrada del agente.

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

Esquema de salida StringResponse (en desuso)

StringResponse permite devolver la respuesta del agente como un objeto con un único campo de cadena content :

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