Schéma de l’agent d’entrée et de sortie hérité

Les agents IA doivent respecter les exigences de schéma d’entrée et de sortie spécifiques pour être compatibles avec d’autres fonctionnalités sur Databricks. Cette page explique comment utiliser les signatures et interfaces de création d’agents héritées : ChatAgent interface, ChatModel interface, SplitChatMessageRequest schéma d’entrée et StringResponse schéma de sortie.

Créer un agent ChatAgent hérité

L’interface MLflow ChatAgent est similaire, mais pas strictement compatible avec le schéma OpenAI ChatCompletion .

Pour découvrir comment créer un ChatAgent, consultez les exemples de la section suivante et la documentation MLflow – Qu’est-ce que l’interface ChatAgent.

Pour créer et déployer des agents à l’aide ChatAgentde , installez les éléments suivants :

• databricks-agents 0.16.0 ou version ultérieure
• mlflow 2.20.2 ou version ultérieure
• Python 3.10 ou version ultérieure.
  • Pour répondre à cette exigence, vous pouvez utiliser le calcul sans serveur ou Databricks Runtime 13.3 LTS ou version ultérieure.

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

Que se passe-t-il si j’ai déjà un agent ?

Si vous disposez déjà d’un agent créé avec LangChain, LangGraph ou une infrastructure similaire, vous n’avez pas besoin de réécrire votre agent pour l’utiliser sur Databricks. Au lieu de cela, il vous suffit d’encapsuler votre agent existant avec l’interface MLflow ChatAgent :

1. Écrire une classe wrapper Python qui hérite de mlflow.pyfunc.ChatAgent.

   Dans la classe wrapper, conservez votre agent existant en tant qu’attribut self.agent = your_existing_agent.

2. La ChatAgent classe vous oblige à implémenter une predict méthode pour gérer les demandes de non-diffusion en continu.

   predict doit accepter :

   • messages: list[ChatAgentMessage], qui est une liste de ChatAgentMessage chacun ayant un rôle (comme « utilisateur » ou « assistant »), l'incitation et un ID.

   • (Facultatif) context: Optional[ChatContext] et custom_inputs: Optional[dict] pour des données supplémentaires.

   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 doit renvoyer 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. Convertir entre les formats

   Dans predict, convertissez les messages list[ChatAgentMessage] entrants dans le format d’entrée attendu par votre agent.

   Une fois que votre agent a généré une réponse, convertissez sa sortie en un ou plusieurs ChatAgentMessage objets et encapsulez-les dans un ChatAgentResponse.

Voici un modèle simplifié pour convertir votre agent :

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())))

Pour obtenir des exemples complets, consultez les blocs-notes de la section suivante.

Exemples ChatAgent

Les notebooks suivants montrent comment créer la diffusion en continu et la non-diffusion en continu ChatAgents à l’aide des bibliothèques populaires OpenAI, LangGraph et AutoGen.

LangGraph

Si vous encapsulez un agent LangChain, vous pouvez utiliser mlflow.langchain.output_parsers.ChatAgentOutputParser pour convertir automatiquement les sorties de LangChain en le schéma MLflow ChatAgentMessage et ChatAgentResponse.

Agent d’appel de l’outil LangGraph

Obtenir un ordinateur portable

OpenAI

Agent d’appel d’outil OpenAI

Obtenir un ordinateur portable

Agent d’appel d’outils de l’API Réponses OpenAI

Obtenir un ordinateur portable

Agent de conversation OpenAI uniquement

Obtenir un ordinateur portable

AutoGen

Agent d'activation d'outil AutoGen

Obtenir un ordinateur portable

DSPy

Agent DSPy uniquement dédié à la conversation

Obtenir un ordinateur portable

Pour savoir comment développer les fonctionnalités de ces agents en ajoutant des outils, consultez outils d’agent IA.

Diffusion en continu des réponses ChatAgent

Les agents de diffusion en continu fournissent des réponses dans un flux continu de segments plus petits et incrémentiels. La diffusion en continu réduit la latence perçue et améliore l’expérience utilisateur pour les agents conversationnels.

Pour créer un streaming ChatAgent, définissez une predict_stream méthode qui retourne un générateur qui génère des ChatAgentChunk objets, chacun ChatAgentChunk contient une partie de la réponse. En savoir plus sur le comportement de diffusion en continu idéal ChatAgent dans la documentation MLflow.

Le code suivant montre un exemple predict_stream de fonction, pour obtenir des exemples complets d’agents de streaming, consultez les exemples 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"]
      )

Concevoir un agent ChatModel ancien

ChatModel est une interface de création d’agent héritée dans MLflow qui étend le schéma ChatCompletion d’OpenAI, ce qui vous permet de maintenir la compatibilité avec les plateformes prenant en charge la norme ChatCompletion tout en ajoutant des fonctionnalités personnalisées. Pour plus d’informations , consultez MLflow : Prise en main de ChatModel .

La création de votre agent en tant que sous-classe de mlflow.pyfunc.ChatModel offre les avantages suivants :

• Active la sortie de l’agent de streaming lors de l'invocation d'un agent servi (en contournant {stream: true} dans le corps de la requête).

• Active automatiquement les tables d’inférence de passerelle IA lorsque votre agent est servi, en fournissant l’accès aux métadonnées de journal de requête améliorées, telles que le nom du demandeur.

  Avertissement

  Les journaux des demandes et les journaux d’évaluation sont déconseillés et seront supprimés dans une version ultérieure. Consultez les journaux des demandes et la dépréciation des journaux d’évaluation pour obtenir des conseils sur la migration.

• Vous permet d’écrire du code de l’agent compatible avec le schéma ChatCompletion à l’aide de classes Python typées.

• MLflow déduit automatiquement une signature compatible avec la finalisation de conversation lors de l'enregistrement de l'agent, même sans un input_example. Cela simplifie le processus d’inscription et de déploiement de l’agent. Consultez Déduire la signature du modèle pendant l’enregistrement.

Le code suivant est le mieux exécuté dans un notebook Databricks. Les notebooks fournissent un environnement pratique pour le développement, le test et l’itération sur votre agent.

La classe MyAgent étend mlflow.pyfunc.ChatModel, implémentant la méthode predict requise. Cela garantit la compatibilité avec Mosaïque AI Agent Framework.

La classe inclut également les méthodes facultatives _create_chat_completion_chunk et les predict_stream pour gérer les sorties de diffusion en continu.

# 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)

Pendant que vous définissez la classe MyAgent d’agent dans un bloc-notes, nous vous recommandons de créer un bloc-notes de pilote distinct. Le bloc-notes du pilote enregistre l’agent dans le Registre de modèles et déploie l’agent à l’aide de Model Serve.

Cette séparation suit le flux de travail recommandé par Databricks pour l'enregistrement des modèles à l'aide de la méthodologie 'Models from Code' de MLflow.

Schéma d’entrée SplitChatMessageRequest (déconseillé)

SplitChatMessagesRequest vous permet de transmettre séparément la requête actuelle et l'historique comme entrée de l'agent.

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

Schéma de sortie StringResponse (déconseillé)

StringResponse vous permet de retourner la réponse de l’agent en tant qu’objet avec un champ de chaîne content unique :

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