Interroger un agent déployé sur Databricks

Découvrez comment envoyer des requêtes aux agents déployés sur Databricks Apps ou aux Endpoint de Model Serving. Databricks propose plusieurs méthodes de query pour s’adapter aux différents cas d’usage 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 OpenAI Databricks (recommandé)	Intégration native, prise en charge complète des fonctionnalités, capacités de streaming
API REST	Compatible OpenAI ; indépendant du langage ; fonctionne avec les outils existants
AI Functions: ai_query	Compatible OpenAI, query uniquement les agents existants hébergés sur les Endpoints de Model Serving

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

Databricks recommande le Client OpenAI de Databricks pour les nouvelles applications. Choisissez l' API REST lors de l'intégration avec des plateformes qui attendent des Endpoint compatibles OpenAI.

Client OpenAI Databricks (Recommandé)

Databricks vous recommande d'utiliser le client DatabricksOpenAI pour query un agent déployé. Selon l'API de votre agent déployé, vous utiliserez soit le client de réponses, soit le client de complétions de discussion :

Agents deployed to Apps

Utilisez l'exemple suivant pour les agents hébergés sur Databricks Apps suivant l'interface ResponsesAgent, qui est l'approche recommandée pour créer des 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 ID de trace de la réponse, incluez l'en-tête x-mlflow-return-trace-id à l'aide de extra_headers. Utilisez ensuite 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 anciens agents hébergés sur Model Serving suivant l'interface ResponsesAgent. Vous pouvez utiliser soit un jeton OAuth Databricks, soit 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 souhaitez transmettre 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 la diffusion de modèles 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 souhaitez transmettre 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 fournit des Endpoint pour les modèles compatibles avec OpenAI. Cela 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 à des plateformes tierces qui s'attendent à des APIs compatibles 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. Consultez la documentation d'authentification de Databricks pour plus d'options et d'informations.

Agents deployed to Apps

Utilisez l'exemple suivant pour les agents hébergés sur Databricks Apps suivant l'interface ResponsesAgent, qui est l'approche recommandée pour créer des 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 à partir de la réponse, incluez l'en-tête x-mlflow-return-trace-id dans votre requête. Le corps de la réponse comprend 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-..."}) vers la fin du flux. 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 Databricks OAuth ou un jeton d'accès personnel (PAT) pour query les agents hébergés sur Model Serving. L'appel API REST est équivalent à :

• 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 ChatAgent ou ChatModel héritées. Cela é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. Voir la fonctionai_query pour la syntaxe SQL et les définitions de paramètres.

SQL

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

Étapes suivantes

Surveiller les applications GenAI en production