Partager via

Interroger un agent Mosaic AI déployé

Découvrez comment envoyer des demandes aux agents déployés sur un point de terminaison Model Serving. Databricks fournit plusieurs méthodes de requête pour répondre à différents cas d’usage et besoins d’intégration.

Pour savoir comment déployer des agents, consultez Déployer un agent pour les applications IA génératives.

Sélectionnez l’approche de requête qui correspond le mieux à votre cas d’usage :

Méthode avantages clés
Client Databricks OpenAI (recommandé) Intégration native, prise en charge complète des fonctionnalités, fonctionnalités de streaming
Déploiements MLflow client Modèles MLflow existants, pipelines ML établis
REST API Compatible avec OpenAI, indépendant du langage, fonctionne avec les outils existants
Fonctions IA : Compatible avec OpenAI, fonctionne avec les outils existants

Databricks recommande le client Databricks OpenAI pour les nouvelles applications. Choisissez l’API REST lors de l’intégration à des plateformes qui attendent des points de terminaison compatibles OpenAI.

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

Points de terminaison de ResponsesAgent

Utilisez l’exemple suivant pour les agents créés avec l’interface ResponsesAgent, qui est l’approche recommandée pour créer des agents.

from databricks.sdk import WorkspaceClient

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

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

## 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 passer custom_inputs ou databricks_options, vous pouvez les ajouter avec le extra_body param :

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)

Points de terminaison ChatAgent ou ChatModel

Utilisez l’exemple suivant pour les agents créés avec les interfaces ChatAgent ou ChatModel héritées, qui sont toujours prises en charge, mais non recommandées pour les nouveaux agents.

from databricks.sdk import WorkspaceClient

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

w = WorkspaceClient()
client = w.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 passer custom_inputs ou databricks_options, vous pouvez les ajouter avec le extra_body param :

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)

Client de déploiement MLflow

Utilisez le client de déploiements MLflow lors de l’utilisation de flux de travail et de pipelines MLflow existants. Cette approche s’intègre naturellement au suivi MLflow et à la gestion des expériences.

Les exemples suivants vous montrent comment interroger un agent à l’aide du client de déploiement MLflow. Pour les nouvelles applications, Databricks recommande d’utiliser databricks OpenAI Client pour ses fonctionnalités améliorées et son intégration native.

En fonction de l’API de votre agent déployé, vous allez utiliser le format ResponsesAgent ou ChatAgent :

Points de terminaison de ResponsesAgent

Utilisez l’exemple suivant pour les agents créés avec l’interface ResponsesAgent, qui est l’approche recommandée pour créer des agents.

from mlflow.deployments import get_deploy_client

client = get_deploy_client()
input_example = {
    "input": [{"role": "user", "content": "What does Databricks do?"}],
    ## Optional: Include any custom inputs
    ## "custom_inputs": {"id": 5},
    "databricks_options": {"return_trace": True},
}
endpoint = "<agent-endpoint-name>" # TODO: update this with your endpoint name

## Call predict for non-streaming responses
response = client.predict(endpoint=endpoint, inputs=input_example)

## Call predict_stream for streaming responses
streaming_response = client.predict_stream(endpoint=endpoint, inputs=input_example)

Points de terminaison ChatAgent ou ChatModel

Utilisez cette option pour les agents créés avec des interfaces ChatAgent ou ChatModel héritées, qui sont toujours prises en charge, mais non recommandées pour les nouveaux agents.

from mlflow.deployments import get_deploy_client

client = get_deploy_client()
input_example = {
    "messages": [{"role": "user", "content": "What does Databricks do?"}],
    ## Optional: Include any custom inputs
    ## "custom_inputs": {"id": 5},
    "databricks_options": {"return_trace": True},
}
endpoint = "<agent-endpoint-name>" # TODO: update this with your endpoint name

## Call predict for non-streaming responses
response = client.predict(endpoint=endpoint, inputs=input_example)

## Call predict_stream for streaming responses
streaming_response = client.predict_stream(endpoint=endpoint, inputs=input_example)

client.predict() et client.predict_stream() appelez les fonctions de l’agent que vous avez définies lors de la création de l’agent. Consultez les réponses de streaming.

REST API

L’API REST Databricks fournit des points de terminaison pour les modèles compatibles Avec OpenAI. Cela vous permet d’utiliser des 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 attendent des API compatibles OpenAI
  • Migration d’OpenAI vers Databricks avec des modifications de code minimales

Authentifiez-vous auprès de l’API REST à l’aide d’un jeton Databricks OAuth ou d’un jeton d’accès personnel (PAT). Les exemples ci-dessous utilisent un jeton OAuth Databricks, reportez-vous à la documentation d’authentification Databricks pour plus d’options et d’informations.

Points de terminaison de ResponsesAgent

Utilisez l’exemple suivant pour les agents créés avec l’interface ResponsesAgent, qui est l’approche recommandée pour créer des agents. L’appel d’API REST équivaut à :

  • Utilisation du client Databricks OpenAI avec responses.create.
  • Envoi d’une requête POST à l’URL du point de terminaison spécifique (par exemple : https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Pour plus d’informations, consultez la page de service du modèle de votre point de terminaison et la documentation sur le service de modèle.
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 passer custom_inputs ou databricks_options, vous pouvez les ajouter avec le extra_body param :

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,
    "extra_body": {
      "custom_inputs": { "id": 5 },
      "databricks_options": { "return_trace": true }
    }
  }'

Points de terminaison ChatAgent ou ChatModel

Utilisez cette option pour les agents créés avec des interfaces ChatAgent ou ChatModel héritées, qui sont toujours prises en charge, mais non recommandées pour les nouveaux agents. Cela équivaut à :

  • Utilisation du client Databricks OpenAI avec chat.completions.create.
  • Envoi d’une requête POST à l’URL du point de terminaison spécifique (par exemple : https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Pour plus d’informations, consultez la page de service du modèle de votre point de terminaison et la documentation sur le service de modèle.
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 passer custom_inputs ou databricks_options, vous pouvez les ajouter avec le extra_body param :

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,
    "extra_body": {
      "custom_inputs": { "id": 5 },
      "databricks_options": { "return_trace": true }
    }
  }'

Fonctions IA : ai_query

Vous pouvez l’utiliser ai_query pour interroger un agent IA déployé à l’aide de SQL. Consultez ai_query la fonction pour la syntaxe SQL et les définitions de paramètres.

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

Étapes suivantes

Surveiller GenAI en production

  • Last updated on