Écrire un agent d'IA et le déployer sur 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.

Cette page explique comment créer un agent d'IA en Python à l'aide de Custom Agents et de bibliothèques populaires de création d'agents telles que LangGraph et OpenAI.

Exigences

astuce

Databricks recommande d'installer la dernière version du client Python MLflow lors du développement d'agents.

Pour créer et déployer des agents en utilisant l'approche décrite sur cette page, installez les éléments suivants :

• databricks-agents 1.2.0 ou supérieur
• mlflow 3.1.3 ou au-dessus
• Python 3.10 ou version ultérieure.
  • Utilisez le compute Serverless ou Databricks Runtime 13.3 LTS ou une version supérieure pour satisfaire à cette exigence.

%pip install -U -qqqq databricks-agents mlflow

Databricks recommande également d'installer les packages d'intégration Databricks AI Bridge pour l'authoring d'agents. Ces packages d'intégration fournissent une couche partagée d'APIs qui interagissent avec les fonctionnalités de Databricks AI, telles que les agents Genie et la recherche AI, à travers les frameworks d'authoring d'agents et les SDK.

OpenAI

%pip install -U -qqqq databricks-openai

LangChain/LangGraph

%pip install -U -qqqq databricks-langchain

DSPy

%pip install -U -qqqq databricks-dspy

Pure Python agents

%pip install -U -qqqq databricks-ai-bridge

Utilisez ResponsesAgent pour créer des agents

Databricks recommande l'interface MLflow ResponsesAgent pour créer des agents de qualité professionnelle. ResponsesAgent vous permet de créer des agents avec n’importe quel framework tiers, puis de l’intégrer aux fonctionnalités de Databricks AI pour des capacités de journalisation, de traçabilité, d’évaluation, de déploiement et de monitoring robustes.

Le schéma ResponsesAgent est compatible avec le schéma OpenAI Responses. Pour en savoir plus sur OpenAI Responses, consultez OpenAI : Responses vs. ChatCompletion.

remarque

L'ancienne interface ChatAgent est toujours prise en charge sur Databricks. Cependant, pour les nouveaux agents, Databricks recommande d’utiliser la dernière version de MLflow et l’interface ResponsesAgent.

Consultez le schéma d'agent d'entrée et de sortie hérité (Model Serving).

ResponsesAgent offre les avantages suivants :

• Capacités d'agent avancées

  • Prise en charge multi-agents
  • Sortie en streaming : diffusez la sortie par petits fragments.
  • Historique complet des messages d'appels d'outils : renvoie plusieurs messages, y compris les messages d'appels d'outils intermédiaires, pour une qualité et une gestion de conversation améliorées.
  • Prise en charge de la confirmation de l'appel d'outil
  • Support des outils à exécution longue

• Développement, déploiement et monitoring simplifiés

  • Créer des agents à l'aide de n'importe quel framework : Encapsulez tout agent existant à l'aide de l'interface ResponsesAgent pour obtenir une compatibilité prête à l'emploi avec AI Playground, Agent Evaluation et le Monitoring d'Agent.
  • Interfaces de création typées : Écrivez du code d'agent en utilisant des classes Python typées, en bénéficiant de la complétion automatique de l'IDE et des Notebooks.
  • Infèrence automatique des signatures : MLflow infère automatiquement les signatures ResponsesAgent lors de la journalisation d'un agent, ce qui simplifie l'enregistrement et le déploiement. Voir Inférer la signature du modèle lors de la journalisation.
  • **Traçage automatique** : MLflow trace predict predict_stream automatiquement vos fonctions et, en agrégeant les réponses en Stream pour une évaluation et un affichage facilités.
  • Tables d'inférence améliorées par AI Gateway : Les tables d'inférence AI Gateway sont automatiquement activées pour les agents déployés, permettant d'accéder aux métadonnées détaillées des logs de requêtes.

Pour savoir comment créer un ResponsesAgent, consultez les exemples de la section suivante et la documentation MLflow - ResponsesAgent pour Model Serving.

ResponsesAgent exemples

Les Notebooks suivants montrent comment créer des ResponsesAgent en streaming et hors streaming à l'aide de bibliothèques populaires. Pour savoir comment étendre les capacités de ces agents, consultez Connecter des agents à des outils.

OpenAI

Agent de chat simple OpenAI utilisant des modèles hébergés par Databricks

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks

Agent d'appel d'outils OpenAI MCP

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks

Agent d'appel d'outils OpenAI utilisant des modèles hébergés par Databricks

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks

Agent d'appel d'outil OpenAI utilisant des modèles hébergés par OpenAI.

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks

LangGraph

Agent d’appel d’outils LangGraph MCP

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks

DSPy

agent DSPy d’appel d’outils à un seul tour

Ouvrir le notebook dans un nouvel ongletOuvrir dans Databricks

Exemple multi-agent

Pour apprendre à créer un système multi-agents, consultez Utiliser Genie dans les systèmes multi-agents (Model Serving).

Exemple d'agent avec état

Pour apprendre à créer des agents avec état dotés d'une mémoire à court et long terme en utilisant Lakebase comme magasin de mémoire, consultez Mémoire d'agent IA (Model Serving).

Exemple d'agent non conversationnel

Contrairement aux agents conversationnels qui gèrent les dialogues à plusieurs tours, les agents non conversationnels se concentrent sur l'exécution efficace de tâches bien définies. Cette architecture rationalisée permet un throughput plus élevé pour les requêtes indépendantes.

Pour savoir comment créer un agent non conversationnel, consultez les agents IA non conversationnels à l'aide de MLflow.

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

Si vous disposez déjà d’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 MLflow ResponsesAgent :

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

   À l'intérieur de la classe d'enveloppe, référencez l'agent existant comme un attribut self.agent = your_existing_agent.

2. La classe ResponsesAgent nécessite l'implémentation d'une méthode predict qui renvoie un ResponsesAgentResponse pour gérer les requêtes non-streaming. Voici un exemple du schéma ResponsesAgentResponses :

   Python

   import uuid
   # input as a dict
   {"input": [{"role": "user", "content": "What did the data scientist say when their Spark job finally completed?"}]}
   
   # output example
   ResponsesAgentResponse(
       output=[
           {
               "type": "message",
               "id": str(uuid.uuid4()),
               "content": [{"type": "output_text", "text": "Well, that really sparked joy!"}],
               "role": "assistant",
           },
       ]
   )

3. Dans la fonction predict, convertissez les messages entrants de ResponsesAgentRequest au format attendu par l'agent. Une fois que l'agent génère une réponse, convertissez sa sortie en un objet ResponsesAgentResponse.

Consultez les exemples de code suivants pour voir comment convertir les agents existants en ResponsesAgent:

Basic conversion

Pour les agents non-streaming, convertissez les entrées et les sorties dans la fonction predict.

Python

from uuid import uuid4

from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
)


class MyWrappedAgent(ResponsesAgent):
    def __init__(self, agent):
        # Reference your existing agent
        self.agent = agent

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        # Convert incoming messages to your agent's format
        # prep_msgs_for_llm is a function you write to convert the incoming messages
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # Call your existing agent (non-streaming)
        agent_response = self.agent.invoke(messages)

        # Convert your agent's output to ResponsesAgent format, assuming agent_response is a str
        output_item = (self.create_text_output_item(text=agent_response, id=str(uuid4())),)

        # Return the response
        return ResponsesAgentResponse(output=[output_item])

Streaming with code re-use

Pour les agents de streaming, vous pouvez faire preuve d'ingéniosité et réutiliser la logique afin d'éviter de dupliquer le code qui convertit les messages :

Python

from typing import Generator
from uuid import uuid4

from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
)


class MyWrappedStreamingAgent(ResponsesAgent):
    def __init__(self, agent):
        # Reference your existing agent
        self.agent = agent

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        """Non-streaming predict: collects all streaming chunks into a single response."""
        # Reuse the streaming logic and collect all output items
        output_items = []
        for stream_event in self.predict_stream(request):
            if stream_event.type == "response.output_item.done":
                output_items.append(stream_event.item)

        # Return all collected items as a single response
        return ResponsesAgentResponse(output=output_items)

    def predict_stream(
        self, request: ResponsesAgentRequest
    ) -> Generator[ResponsesAgentStreamEvent, None, None]:
        """Streaming predict: the core logic that both methods use."""
        # Convert incoming messages to your agent's format
        # prep_msgs_for_llm is a function you write to convert the incoming messages, included in full examples linked below
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # Stream from your existing agent
        item_id = str(uuid4())
        aggregated_stream = ""
        for chunk in self.agent.stream(messages):
            # Convert each chunk to ResponsesAgent format
            yield self.create_text_delta(delta=chunk, item_id=item_id)
            aggregated_stream += chunk

        # Emit an aggregated output_item for all the text deltas with id=item_id
        yield ResponsesAgentStreamEvent(
            type="response.output_item.done",
            item=self.create_text_output_item(text=aggregated_stream, id=item_id),
        )

Migrate from ChatCompletions

Si votre agent existant utilise l'API ChatCompletions OpenAI, vous pouvez le migrer vers ResponsesAgent sans réécrire sa logique de base. Ajouter un wrapper qui :

1. Convertit les messages ResponsesAgentRequest entrants au format ChatCompletions attendu par votre agent.
2. Traduit les sorties ChatCompletions dans le schéma ResponsesAgentResponse.
3. Prend en charge le streaming en option en mappant les deltas incrémentiels de ChatCompletions en objets ResponsesAgentStreamEvent.

Python

from typing import Generator
from uuid import uuid4

from databricks.sdk import WorkspaceClient
from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
)


# Legacy agent that outputs ChatCompletions objects
class LegacyAgent:
    def __init__(self):
        self.w = WorkspaceClient()
        self.OpenAI = self.w.serving_endpoints.get_open_ai_client()

    def stream(self, messages):
        for chunk in self.OpenAI.chat.completions.create(
            model="databricks-claude-sonnet-4-5",
            messages=messages,
            stream=True,
        ):
            yield chunk.to_dict()


# Wrapper that converts the legacy agent to a ResponsesAgent
class MyWrappedStreamingAgent(ResponsesAgent):
    def __init__(self, agent):
        # `agent` is your existing ChatCompletions agent
        self.agent = agent

    def prep_msgs_for_llm(self, messages):
        # dummy example of prep_msgs_for_llm
        # real example of prep_msgs_for_llm included in full examples linked below
        return [{"role": "user", "content": "Hello, how are you?"}]

    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        """Non-streaming predict: collects all streaming chunks into a single response."""
        # Reuse the streaming logic and collect all output items
        output_items = []
        for stream_event in self.predict_stream(request):
            if stream_event.type == "response.output_item.done":
                output_items.append(stream_event.item)

        # Return all collected items as a single response
        return ResponsesAgentResponse(output=output_items)

    def predict_stream(
        self, request: ResponsesAgentRequest
    ) -> Generator[ResponsesAgentStreamEvent, None, None]:
        """Streaming predict: the core logic that both methods use."""
        # Convert incoming messages to your agent's format
        messages = self.prep_msgs_for_llm([i.model_dump() for i in request.input])

        # process the ChatCompletion output stream
        agent_content = ""
        tool_calls = []
        msg_id = None
        for chunk in self.agent.stream(messages):  # call the underlying agent's stream method
            delta = chunk["choices"][0]["delta"]
            msg_id = chunk.get("id", None)
            content = delta.get("content", None)
            if tc := delta.get("tool_calls"):
                if not tool_calls:  # only accommodate for single tool call right now
                    tool_calls = tc
                else:
                    tool_calls[0]["function"]["arguments"] += tc[0]["function"]["arguments"]
            elif content is not None:
                agent_content += content
                yield ResponsesAgentStreamEvent(**self.create_text_delta(content, item_id=msg_id))

        # aggregate the streamed text content
        yield ResponsesAgentStreamEvent(
            type="response.output_item.done",
            item=self.create_text_output_item(agent_content, msg_id),
        )

        for tool_call in tool_calls:
            yield ResponsesAgentStreamEvent(
                type="response.output_item.done",
                item=self.create_function_call_item(
                    str(uuid4()),
                    tool_call["id"],
                    tool_call["function"]["name"],
                    tool_call["function"]["arguments"],
                ),
            )


agent = MyWrappedStreamingAgent(LegacyAgent())

for chunk in agent.predict_stream(
    ResponsesAgentRequest(input=[{"role": "user", "content": "Hello, how are you?"}])
):
    print(chunk)

Pour des exemples complets, voir ResponsesAgent exemples.

Réponses en streaming

Le streaming permet aux agents d'envoyer des réponses par blocs en temps réel au lieu d'attendre la réponse complète. Pour implémenter le streaming avec ResponsesAgent, émettez une série d'événements delta suivie d'un événement d'achèvement final :

1. Émettre des événements delta : Envoyez plusieurs Stream avec le output_text.delta même item_id pour diffuser des blocs de texte en temps réel.
2. Terminer avec l'événement terminé : Envoyez un événement final response.output_item.done avec le même item_id que les événements delta contenant le texte de sortie final complet.

Chaque événement delta Stream un segment de texte au client. L’événement « done » final contient le texte de réponse complet et signale à Databricks de faire ce qui suit :

• Tracez la sortie de votre agent avec le traçage MLflow
• Agréger les réponses Stream dans les tables d'inférence AI Gateway
• Afficher le résultat complet dans l’interface utilisateur d’AI Playground

Propagation des erreurs en streaming

Databricks propage toutes les erreurs rencontrées lors du streaming avec le dernier jeton sous databricks_output.error. Il appartient au client appelant de gérer et de faire remonter correctement cette erreur.

Bash

{
  "delta": …,
  "databricks_output": {
    "trace": {...},
    "error": {
      "error_code": BAD_REQUEST,
      "message": "TimeoutException: Tool XYZ failed to execute."
    }
  }
}

Fonctionnalités avancées

Entrées et sorties personnalisées

Certains scénarios peuvent nécessiter des entrées d'agent supplémentaires, telles que client_type et session_id, ou des sorties comme des liens vers la source de récupération qui ne devraient pas être inclus dans l'historique de chat pour les interactions futures.

Pour ces scénarios, MLflow ResponsesAgent prend en charge en mode natif les champs custom_inputs et custom_outputs. Vous pouvez accéder aux entrées personnalisées via request.custom_inputs dans tous les exemples liés ci-dessus dans Exemples de l'Agent de Réponses.

attention

L'application de révision Agent Evaluation ne prend pas en charge le rendu des traces pour les agents avec des champs de saisie supplémentaires.

Consultez les Notebooks suivants pour apprendre à définir des entrées et des sorties personnalisées.

Fournissez custom_inputs dans l'AI Playground et l'application d'évaluation

Si votre agent accepte des entrées supplémentaires à l'aide du champ custom_inputs, vous pouvez fournir manuellement ces entrées à la fois dans l'AI Playground et dans l'application de révision.

1. Dans l'AI Playground ou l'application Agent Review, sélectionnez l'icône d'engrenage .

2. Activer custom_inputs .

3. Fournissez un objet JSON qui correspond au schéma d'entrée défini de votre agent.

   

Spécifier des schémas de récupérateur personnalisés

Les agents d'IA utilisent couramment des récupérateurs pour trouver et interroger des données non structurées à partir des index de recherche IA. Par exemple, pour les outils de récupération, voir Connecter des agents à des données non structurées.

Tracez ces retrievers au sein de votre agent avec les spans RETRIEVER de MLflow pour activer les fonctionnalités produit de Databricks, notamment :

• Affichage automatique des liens vers les documents source récupérés dans l'interface utilisateur de l'AI Playground
• Exécution automatique des juges d'ancrage de récupération et de pertinence dans Agent Evaluation

remarque

Databricks recommande d'utiliser les outils de récupération fournis par les packages Databricks AI Bridge comme databricks_langchain.VectorSearchRetrieverTool et databricks_openai.VectorSearchRetrieverTool, car ils sont déjà conformes au schéma de récupération MLflow. Consultez Développer un outil de récupération localement à l'aide du pont Databricks AI.

Si votre agent inclut des portées de récupérateur avec un schéma personnalisé, appelez mlflow.models.set_retriever_schema lorsque vous définissez votre agent dans le code. Ceci mappe les colonnes de sortie de votre récupérateur aux champs attendus de MLflow (primary_key, text_column, doc_uri).

Python

import mlflow
# Define the retriever's schema by providing your column names
# For example, the following call specifies the schema of a retriever that returns a list of objects like
# [
#     {
#         'document_id': '9a8292da3a9d4005a988bf0bfdd0024c',
#         'chunk_text': 'MLflow is the largest open source AI engineering platform for agents, LLMs, and ML models...',
#         'doc_uri': 'https://mlflow.org/docs/latest/index.html',
#         'title': 'MLflow: The Largest Open Source AI Engineering Platform'
#     },
#     {
#         'document_id': '7537fe93c97f4fdb9867412e9c1f9e5b',
#         'chunk_text': 'A great way to get started with MLflow is to use the autologging feature. Autologging automatically logs your model...',
#         'doc_uri': 'https://mlflow.org/docs/latest/getting-started/',
#         'title': 'Getting Started with MLflow'
#     },
# ...
# ]
mlflow.models.set_retriever_schema(
    # Specify the name of your retriever span
    name="mlflow_docs_vector_search",
    # Specify the output column name to treat as the primary key (ID) of each retrieved document
    primary_key="document_id",
    # Specify the output column name to treat as the text content (page content) of each retrieved document
    text_column="chunk_text",
    # Specify the output column name to treat as the document URI of each retrieved document
    doc_uri="doc_uri",
    # Specify any other columns returned by the retriever
    other_columns=["title"],
)

remarque

La colonne doc_uri est particulièrement importante lors de l'évaluation des performances du récupérateur. doc_uri est l'identifiant principal des documents renvoyés par le récupérateur, vous permettant de les comparer aux ensembles d'évaluation de vérité terrain. Consultez Ensembles d'évaluation (MLflow 2).

Considérations de déploiement

Préparer Databricks Model Serving

Databricks déploie des ResponsesAgentdans un environnement distribué sur Databricks Model Serving. Cela signifie que, lors d’une conversation multi-tours, la même réplique de service pourrait ne pas traiter toutes les requêtes. Veuillez prêter attention aux implications suivantes pour la gestion de l’état des agents :

• Évitez la mise en cache locale : lors du déploiement ResponsesAgent d'un, ne partez pas du principe que la même réplique gère toutes les requêtes dans une conversation à plusieurs tours. Reconstruire l'état interne à l'aide d'un dictionnaire ResponsesAgentRequest schéma pour chaque tour.

• **État thread-safe** : Concevez l'état de l'agent pour qu'il soit thread-safe, prévenant ainsi les conflits dans les environnements multithreads.

• Initialisez l'état dans la fonction predict : initialisez l'état chaque fois que la fonction predict est appelée, pas pendant l'initialisation de ResponsesAgent. Le stockage de l'état au niveau ResponsesAgent pourrait entraîner une fuite d'informations entre les conversations et provoquer des conflits, car une seule réplique de ResponsesAgent pourrait gérer les requêtes de plusieurs conversations.

Paramétrer le code pour le déploiement entre les environnements

Paramétrez le code d'agent pour réutiliser le même code d'agent dans différents environnements.

Les parameters sont des paires clé-valeur que vous définissez dans un dictionnaire Python ou un fichier .yaml.

Pour configurer le code, créez un(e) ModelConfig en utilisant un dictionnaire Python ou un fichier .yaml. ModelConfig est un ensemble de parameters clé-valeur qui permet une gestion flexible de la configuration. Par exemple, vous pouvez utiliser un dictionnaire pendant le développement, puis le convertir en un fichier .yaml pour le déploiement en production et la CI/CD.

Un exemple de ModelConfig est présenté ci-dessous :

YAML

llm_parameters:
  max_tokens: 500
  temperature: 0.01
model_serving_endpoint: databricks-meta-llama-3-3-70b-instruct
vector_search_index: ml.docs.databricks_docs_index
prompt_template: 'You are a hello world bot. Respond with a reply to the user''s
  question that indicates your prompt template came from a YAML file. Your response
  must use the word "YAML" somewhere. User''s question: {question}'
prompt_template_input_vars:
  - question

Dans votre code d'agent, vous pouvez référencer une configuration default (de développement) à partir du fichier ou du dictionnaire .yaml :

Python

import mlflow
# Example for loading from a .yml file
config_file = "configs/hello_world_config.yml"
model_config = mlflow.models.ModelConfig(development_config=config_file)

# Example of using a dictionary
config_dict = {
    "prompt_template": "You are a hello world bot. Respond with a reply to the user's question that is fun and interesting to the user. User's question: {question}",
    "prompt_template_input_vars": ["question"],
    "model_serving_endpoint": "databricks-meta-llama-3-3-70b-instruct",
    "llm_parameters": {"temperature": 0.01, "max_tokens": 500},
}

model_config = mlflow.models.ModelConfig(development_config=config_dict)

# Use model_config.get() to retrieve a parameter value
# You can also use model_config.to_dict() to convert the loaded config object
# into a dictionary
value = model_config.get('sample_param')

Ensuite, lors de l'enregistrement de votre agent, spécifiez le parameter model_config à log_model pour spécifier un ensemble de parameters personnalisé à utiliser lors du chargement de l'agent enregistré. Consultez documentation MLflow – ModelConfig.

Utilisez un code synchrone ou des modèles de rappel

Pour garantir la stabilité et la compatibilité, utilisez un code synchrone ou des modèles basés sur des rappels dans votre implémentation d'agent.

Databricks gère automatiquement la communication asynchrone pour fournir une concurrence et des performances optimales lorsque vous déployez un agent. L'introduction de boucles d'événements personnalisées ou de frameworks asynchrones pourrait entraîner des erreurs comme RuntimeError: This event loop is already running and caused unpredictable behavior.

Databricks recommande d'éviter la programmation asynchrone, telle que l'utilisation d'asyncio ou la création de boucles d'événements personnalisées, lors du développement d'agents.

Ressources supplémentaires

• Créez vos propres outils d'agent.
• Enregistrer un agent IA.
• Ajouter des traces à un agent d'IA.
• Déployer un agent IA.