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 des agents, 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 adhérer à des 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 les interfaces de création d'agents héritées : l'interface ChatAgent, l'interface ChatModel, le schéma d'entrée SplitChatMessageRequest et le schéma de sortie StringResponse.

Créez un agent ChatAgent hérité​

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

ChatAgent intègre facilement les agents existants pour la compatibilité Databricks.

Pour apprendre à 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 supérieur
  • 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 vous avez déjà un agent ?​

Si vous avez déjà un agent construit 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 ChatAgent MLflow :

  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 classe ChatAgent vous demande d'implémenter une predict méthode 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 (tel que « user » 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.

    Une fois que votre agent a généré 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 intégrez 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.

Voici un Template simplifié pour convertir 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 non en streaming à l'aide des bibliothèques populaires OpenAI, LangGraph et AutoGen.

Si vous intégrez 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 LangGraph d'appel d'outils

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks

Pour apprendre à étendre les capacités de ces agents en ajoutant des outils, veuillez consulter Connecter des agents aux outils.

Réponses de streaming ChatAgent​

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

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

Le code suivant présente 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 créer des agents ou des applications d'IA générative. Pour migrer de ChatModel à ChatAgent, consultez la documentation MLflow : Migrer de ChatModel à ChatAgent.

ChatModel est une interface d'agent de création héritée dans MLflow qui étend le schéma ChatCompletion d'OpenAI, vous permettant de maintenir la compatibilité avec les plateformes prenant en charge la norme ChatCompletion tout en ajoutant des fonctionnalités personnalisées. Consultez MLflow : Pour commencer avec ChatModel pour en savoir plus.

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

  • Permet la sortie de l'agent en 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 la passerelle IA lorsque votre agent est déployé, fournissant 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 future version. Consultez la dépréciation des logs de requête et des logs d'évaluation pour obtenir des conseils de migration.

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

  • MLflow déduit automatiquement une signature compatible avec la complétion de chat lors de l'enregistrement 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 lors de l'enregistrement.

Le code suivant est préférable d'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. Ceci garantit la compatibilité avec les agents personnalisés.

La classe inclut également les méthodes facultatives _create_chat_completion_chunk et predict_stream pour gérer les sorties de 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)

Bien que vous définissiez la classe d'agent MyAgent dans un Notebook, nous vous recommandons de créer un Notebook de Driver distinct. Le Notebook Driver enregistre l'agent dans le Model Registry et déploie l'agent via le Model Serving.

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

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

SplitChatMessagesRequest vous permet de passer la query actuelle et l'historique séparément comme entrée de l'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 en tant qu'objet avec un champ content de chaîne unique :

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