Aller au contenu principal
Dernière mise à jour le

Schéma d'agent d'entrée et de sortie hérité (Model Serving)

info

Pour les nouveaux cas d’utilisation, Databricks recommande de déployer des agents sur Databricks Apps pour un contrôle total sur le code d’agent, la configuration du serveur et le workflow de déploiement. Consultez Créer un agent IA et le déployer sur Databricks Apps. Pour migrer un agent existant, consultez Migrer un agent de Model Serving vers Databricks Apps.

remarque

Databricks recommande de migrer vers le schéma ResponsesAgent pour créer des agents. Consultez Créer un agent IA et le déployer sur Databricks Apps.

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

Créer un agent ChatAgent hérité​

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

ChatAgent enveloppe facilement les agents existants pour la compatibilité Databricks.

Pour savoir 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 de ChatAgent, installez les éléments suivants :

  • databricks-agents0.16.0 ou supérieur
  • mlflow 2.20.2 ou version ultérieure
  • Python 3.10 ou version ultérieure.
    • Pour répondre à cette exigence, vous pouvez utiliser le compute Serverless ou Databricks Runtime 13.3 LTS ou une version ultérieure.
Python
%pip install -U -qqqq databricks-agents==0.16.0 mlflow==2.20.2

Et si j'ai déjà un agent ?​

Si vous avez déjà un agent créé avec LangChain, LangGraph ou un framework similaire, vous n'avez pas besoin de réécrire votre agent pour l'utiliser sur Databricks. Au lieu de cela, enveloppez simplement votre agent existant avec l'interface MLflow ChatAgent :

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

    Dans la classe wrapper, conservez votre agent existant comme attribut self.agent = your_existing_agent.

  2. La classe ChatAgent vous demande d'implémenter une méthode predict pour gérer les requêtes non-streaming.

    predict doit accepter :

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

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

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

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

    Après que votre agent génère une réponse, convertissez sa sortie en un ou plusieurs objets ChatAgentMessage et enveloppez-les dans un ChatAgentResponse.

astuce

Convertir automatiquement la sortie LangChain

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

Le suivant est un Template simplifié pour la conversion de votre agent :

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

Pour des exemples complets, consultez les notebooks dans la section suivante.

ChatAgent exemples​

Les notebooks suivants montrent comment créer des ChatAgents en streaming et hors streaming en utilisant les bibliothèques populaires OpenAI, LangGraph et AutoGen.

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

Agent d'appel d'outils LangGraph

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks

Pour savoir comment étendre les capacités de ces agents en ajoutant des outils, consultez Connecter des agents aux outils.

Réponses de Streaming ChatAgent​

Les agents de streaming fournissent des réponses dans un flux continu de fragments plus petits et incrémentaux. Le streaming réduit la latence perçue et améliore l'expérience utilisateur des agents conversationnels.

Pour créer un ChatAgent de streaming, définissez une méthode predict_stream qui renvoie un générateur qui produit ChatAgentChunk objets - chaque ChatAgentChunk contient une partie de la réponse. En savoir plus sur le comportement de streaming ChatAgent idéal dans la documentation MLflow.

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

Créer un agent ChatModel hérité​

important

Databricks recommande l'interface ChatAgent pour la création d'agents ou d'applications d'IA générative. Pour migrer de ChatModel vers ChatAgent, consultez la documentation MLflow - Migration de ChatModel vers ChatAgent.

ChatModel est une ancienne interface de création d'agents dans MLflow qui étend le schéma ChatCompletion d'OpenAI, vous permettant de maintenir la compatibilité avec les plateformes prenant en charge le standard ChatCompletion tout en ajoutant des fonctionnalités personnalisées. Consultez MLflow : Premiers pas avec ChatModel pour plus de détails.

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

  • Permet la sortie d'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 AI Gateway lorsque votre agent est servi, offrant un accès à des métadonnées de Logs de requêtes améliorées, telles que le nom du demandeur.

attention

Les Logs de requête et les Logs d'évaluation sont obsolètes et seront supprimés dans une prochaine version. Consultez la dépréciation des Logs de requête et des Logs d’évaluation pour obtenir des conseils sur la migration.

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

  • MLflow infère automatiquement une signature compatible avec la complétion de chat lors de la journalisation de l'agent, même sans input_example. Cela simplifie le processus d'enregistrement et de déploiement de l'agent. Consultez Inférer la signature du modèle pendant l'enregistrement.

Le code suivant est préférable de l'exécuter dans un Notebook Databricks. Les Notebooks offrent un environnement pratique pour développer, tester et itérer sur votre agent.

La classe MyAgent étend mlflow.pyfunc.ChatModel, implémentant la méthode predict requise. Cela assure la compatibilité avec les Agents personnalisés.

La classe comprend également les méthodes facultatives _create_chat_completion_chunk et predict_stream pour gérer les sorties en streaming.

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)

Tandis que vous définissez la classe d'agent MyAgent dans un Notebook, nous recommandons de créer un Notebook Driver séparé. Le notebook du Driver enregistre l'agent dans le Model Registry et déploie l'agent à l'aide de Model Serving.

Cette séparation suit le workflow recommandé par Databricks pour la journalisation des modèles en utilisant la méthodologie Models from Code de MLflow.

Schéma d'entrée de SplitChatMessageRequest (obsolète)​

SplitChatMessagesRequest vous permet de transmettre la query et l'historique actuels séparément en tant qu'entrée d'agent.

Python
  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éprécié)​

StringResponse vous permet de renvoyer la réponse de l'agent sous la forme d'un objet avec un champ de chaîne unique content :

{"content": "This is an example string response"}
Sur cette page