Interroger un agent déployé sur Databricks

Découvrez comment envoyer des requêtes aux agents déployés sur les Databricks Apps ou les Endpoint Model Serving. Databricks propose plusieurs méthodes de query pour s’adapter à différents cas d’utilisation et besoins d’intégration.

Sélectionnez l'approche de query qui correspond le mieux à votre cas d'utilisation :

Méthode	Avantages clés
Client Databricks OpenAI (recommandé)	Intégration native, prise en charge complète des fonctionnalités, capacités de streaming
API REST	Compatible avec OpenAI, indépendant du langage, fonctionne avec les outils existants
AI Functions: ai_query	Compatible OpenAI, pour query les agents hérités hébergés uniquement sur les Endpoint de Model Serving

Méthode	Avantages clés
Client Databricks OpenAI (recommandé)	Intégration native, prise en charge complète des fonctionnalités, capacités de streaming
API REST	Compatible avec OpenAI, indépendant du langage, fonctionne avec les outils existants
AI Functions: ai_query	Compatible OpenAI, pour query les agents hérités hébergés uniquement sur les Endpoint de Model Serving

Databricks recommande le **client OpenAI Databricks** pour les nouvelles applications. Choisissez l'**API REST** lors de l'intégration avec des plateformes qui s'attendent à des endpoints compatibles OpenAI.

Client Databricks OpenAI (Recommandé)

Databricks vous recommande d'utiliser le client DatabricksOpenAI pour interroger un agent déployé. Selon l'API de votre agent déployé, vous utiliserez les réponses ou le client de complétions de chat :

Agents deployed to Apps

Utilisez l'exemple suivant pour les agents hébergés sur Databricks Apps suivant l'interface ResponsesAgent, ce qui est l'approche recommandée pour la création d'agents. Vous devez utiliser un jeton OAuth Databricks pour interroger les agents hébergés sur Databricks Apps.

Python

from databricks.sdk import WorkspaceClient
from databricks_openai import DatabricksOpenAI

input_msgs = [{"role": "user", "content": "What does Databricks do?"}]
app_name = "<agent-app-name>"  # TODO: update this with your app name

# The WorkspaceClient must be configured with OAuth authentication
# See: https://docs.databricks.com/aws/en/dev-tools/auth/oauth-u2m.html
w = WorkspaceClient()

client = DatabricksOpenAI(workspace_client=w)

# Run for non-streaming responses. Calls the "invoke" method
# Include the "apps/" prefix in the model name
response = client.responses.create(model=f"apps/{app_name}", input=input_msgs)
print(response)

# Include stream=True for streaming responses. Calls the "stream" method
# Include the "apps/" prefix in the model name
streaming_response = client.responses.create(
    model=f"apps/{app_name}", input=input_msgs, stream=True
)
for chunk in streaming_response:
    print(chunk)

Si vous souhaitez transmettre custom_inputs, vous pouvez les ajouter avec le paramètre extra_body :

Python

streaming_response = client.responses.create(
    model=f"apps/{app_name}",
    input=input_msgs,
    stream=True,
    extra_body={
        "custom_inputs": {"id": 5},
    },
)
for chunk in streaming_response:
    print(chunk)

Pour récupérer un identifiant de trace à partir de la réponse, incluez l'en-tête x-mlflow-return-trace-id à l'aide de extra_headers. Ensuite, utilisez MLflow get_trace pour récupérer la trace complète.

Python

response = client.responses.create(
    model=f"apps/{app_name}",
    input=input_msgs,
    extra_headers={"x-mlflow-return-trace-id": "true"},
)
trace_id = response.metadata["trace_id"]
trace = client.get_trace(trace_id)

Agents on Model Serving

Utilisez l'exemple suivant pour les agents hérités hébergés sur Model Serving suivant l'interface ResponsesAgent. Vous pouvez utiliser un jeton OAuth Databricks ou un jeton d’accès personnel (PAT) pour interroger les agents hébergés sur Model Serving.

Python

from databricks_openai import DatabricksOpenAI

input_msgs = [{"role": "user", "content": "What does Databricks do?"}]
endpoint = "<agent-endpoint-name>" # TODO: update this with your endpoint name

client = DatabricksOpenAI()

# Run for non-streaming responses. Invokes `predict`
response = client.responses.create(model=endpoint, input=input_msgs)
print(response)

# Include stream=True for streaming responses. Invokes `predict_stream`
streaming_response = client.responses.create(model=endpoint, input=input_msgs, stream=True)
for chunk in streaming_response:
  print(chunk)

Si vous voulez passer custom_inputs ou databricks_options, vous pouvez les ajouter avec le paramètre extra_body :

Python

streaming_response = client.responses.create(
    model=endpoint,
    input=input_msgs,
    stream=True,
    extra_body={
        "custom_inputs": {"id": 5},
        "databricks_options": {"return_trace": True},
    },
)
for chunk in streaming_response:
    print(chunk)

Utilisez l’exemple suivant pour les agents hérités sur le service de modèle, en suivant les interfaces ChatAgent ou ChatModel.

Python

from databricks.sdk import WorkspaceClient

messages = [{"role": "user", "content": "What does Databricks do?"}]
endpoint = "<agent-endpoint-name>" # TODO: update this with your endpoint name

ws_client = WorkspaceClient()
client = ws_client.serving_endpoints.get_open_ai_client()

# Run for non-streaming responses. Invokes `predict`
response = client.chat.completions.create(model=endpoint, messages=messages)
print(response)

# Include stream=True for streaming responses. Invokes `predict_stream`
streaming_response = client.chat.completions.create(model=endpoint, messages=messages, stream=True)
for chunk in streaming_response:
  print(chunk)

Si vous voulez passer custom_inputs ou databricks_options, vous pouvez les ajouter avec le paramètre extra_body :

Python

streaming_response = client.chat.completions.create(
    model=endpoint,
    messages=messages,
    stream=True,
    extra_body={
        "custom_inputs": {"id": 5},
        "databricks_options": {"return_trace": True},
    },
)
for chunk in streaming_response:
    print(chunk)

API REST

L'API REST Databricks propose des Endpoints pour les modèles compatibles avec OpenAI. Ceci vous permet d'utiliser les agents Databricks pour servir des applications qui nécessitent des interfaces OpenAI.

Cette approche est idéale pour :

• Applications indépendantes du langage qui utilisent des requêtes HTTP
• Intégration avec des plateformes tierces qui attendent des APIs compatibles avec OpenAI.
• Migration d'OpenAI vers Databricks avec un minimum de modifications de code

Authentifiez-vous auprès de l’API REST à l’aide d’un jeton OAuth Databricks. Reportez-vous à la documentation d’authentification Databricks pour plus d’options et d’information.

Agents deployed to Apps

Utilisez l'exemple suivant pour les agents hébergés sur Databricks Apps suivant l'interface ResponsesAgent, ce qui est l'approche recommandée pour la création d'agents. Vous devez utiliser un jeton OAuth Databricks pour interroger les agents hébergés sur Databricks Apps.

Bash

curl --request POST \
  --url <app-url>.databricksapps.com/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "input": [{ "role": "user", "content": "hi" }],
    "stream": true
  }'

Si vous souhaitez transmettre custom_inputs, vous pouvez les ajouter au corps de la requête :

Bash

curl --request POST \
  --url <app-url>.databricksapps.com/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "input": [{ "role": "user", "content": "hi" }],
    "stream": true,
    "custom_inputs": { "id": 5 }
  }'

Pour récupérer un ID de trace de la réponse, incluez l'en-tête x-mlflow-return-trace-id dans votre requête. Le corps de la réponse inclut un champ metadata.trace_id contenant l'ID de trace. Pour les requêtes de streaming, l'ID de trace est envoyé comme un événement SSE distinct (data: {"trace_id": "tr-..."}) près de la fin du Stream. Utilisez ensuite MLflow get_trace pour récupérer la trace complète à l'aide de l'ID de trace.

Bash

curl --request POST \
  --url <app-url>.databricksapps.com/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --header 'x-mlflow-return-trace-id: true' \
  --data '{
    "input": [{ "role": "user", "content": "hi" }]
  }'

Agents on Model Serving

Utilisez l'exemple suivant pour les agents hérités hébergés sur Model Serving suivant l'interface ResponsesAgent. Vous pouvez utiliser un jeton OAuth Databricks ou un Jeton d'accès personnel (PAT) pour query les agents hébergés sur Model Serving. L'appel de l'API REST équivaut à :

• Utilisation du client Databricks OpenAI avec responses.create.
• Envoi d'une requête POST à l'URL de l'Endpoint spécifique (ex : https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Pour plus d'informations, consultez la page Model Serving de votre endpoint et la documentation Model Serving.

Bash

curl --request POST \
  --url https://<host.databricks.com\>/serving-endpoints/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "model": "\<model-name\>",
    "input": [{ "role": "user", "content": "hi" }],
    "stream": true
  }'

Si vous souhaitez transmettre custom_inputs ou databricks_options, vous pouvez les ajouter au corps de la requête :

Bash

curl --request POST \
  --url https://<host.databricks.com\>/serving-endpoints/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "model": "\<model-name\>",
    "input": [{ "role": "user", "content": "hi" }],
    "stream": true,
    "custom_inputs": { "id": 5 },
    "databricks_options": { "return_trace": true }
  }'

Utilisez ce qui suit pour les agents créés avec les interfaces héritées ChatAgent ou ChatModel. Ceci équivaut à :

• Utilisation du client Databricks OpenAI avec chat.completions.create.
• Envoi d'une requête POST à l'URL de l'Endpoint spécifique (ex : https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Pour plus d'informations, consultez la page Model Serving de votre endpoint et la documentation Model Serving.

Bash

curl --request POST \
  --url https://<host.databricks.com\>/serving-endpoints/chat/completions \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "model": "\<model-name\>",
    "messages": [{ "role": "user", "content": "hi" }],
    "stream": true
  }'

Si vous souhaitez transmettre custom_inputs ou databricks_options, vous pouvez les ajouter au corps de la requête :

Bash

curl --request POST \
  --url https://<host.databricks.com\>/serving-endpoints/chat/completions \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "model": "\<model-name\>",
    "messages": [{ "role": "user", "content": "hi" }],
    "stream": true,
    "custom_inputs": { "id": 5 },
    "databricks_options": { "return_trace": true }
  }'

AI Functions: ai_query

Vous pouvez utiliser ai_query pour interroger un agent déployé hébergé sur Model Serving à l'aide de SQL. Consultez la fonctionai_query pour la syntaxe SQL et les définitions de parameter.

SQL

SELECT ai_query(
  "<model name>", question
) FROM (VALUES ('what is MLflow?'), ('how does MLflow work?')) AS t(question);

Étapes suivantes

Surveillez les applications GenAI en production