Esquema del agente de entrada y salida heredados

Los agentes de inteligencia artificial 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 usar las interfaces e firmas de creación de agentes heredados: ChatAgent interfaz, ChatModel interfaz, SplitChatMessageRequest esquema de entrada y StringResponse esquema de salida.

Creación de un agente de 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, simplemente encapsula 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 indicación y un identificador.

   • (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értalos en uno o varios ChatAgentMessage objetos y envuélvalos en un ChatAgentResponse.

A continuación se muestra una 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 streaming y no streaming ChatAgents mediante las bibliotecas populares 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 cuaderno

OpenAI

Agente de llamada a herramientas de OpenAI

Obtener el cuaderno

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

Obtener el cuaderno

Agente de solo chat de OpenAI

Obtener el cuaderno

AutoGen

Agente de invocación de herramientas de AutoGen

Obtener el cuaderno

DSPy

Agente de solo chat de DSPy

Obtener el cuaderno

Para obtener información sobre cómo expandir las funcionalidades de estos agentes mediante la adición de herramientas, consulte herramientas del agente de IA.

Respuestas de ChatAgent de streaming

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 heredado de ChatModel

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:

• Habilita la salida del agente de streaming al invocar un agente servido (pasando {stream: true} 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.

  Advertencia

  Los registros de solicitud y los registros de evaluación están en desuso y se quitarán en una versión futura. Consulte los registros de peticiones y la desaprobación de los registros de evaluación para obtener instrucciones de 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 .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 proporcionan un entorno conveniente para desarrollar, probar e iterar en el agente.

La clase MyAgent extiende mlflow.pyfunc.ChatModel, implementando el método predict necesario. Esto garantiza la compatibilidad con mosaic AI Agent Framework.

La clase también incluye los métodos opcionales _create_chat_completion_chunk y predict_stream para controlar las salidas 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)

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 te permite enviar por separado la consulta y el historial actuales como entradas para el 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"}