Aller au contenu principal

Interroger les services de modèle

info

Bêta

Cette fonctionnalité est en Bêta. Les administrateurs de compte peuvent contrôler l'accès à cette fonctionnalité depuis la page Aperçus de la console de compte. Consultez Gérer les aperçus Databricks.

Cette page décrit comment query les services de modèle dans Unity Catalog en utilisant les APIs prises en charge.

Exigences

APIs et intégrations prises en charge

Unity AI Gateway prend en charge les APIs et intégrations suivantes :

Query model services with ai_query

Vous pouvez utiliser la fonction ai_query pour query les services de modèle fournis par Databricks directement depuis SQL ou Python. Cela vous permet de capturer les informations de suivi d'utilisation pour vos charges de travail d'inférence par batch.

remarque
  • ai_query la prise en charge de Unity AI Gateway n'est disponible que pour les services de modèle fournis par Databricks (par exemple, databricks-gpt-5-4 ou databricks-claude-sonnet-4). Les services de modèles que vous créez dans Unity AI Gateway ne sont pas encore pris en charge.
  • Seul le suivi de l'utilisation s'applique aux ai_query charges de travail d'inférence par batch. Les autres fonctionnalités de Unity AI Gateway, telles que les limites de taux, les garde-fous, les tables d'inférence et les fallback, ne s'appliquent pas.

Pour commencer :

  1. Activez l'aperçu de Unity AI Gateway pour votre compte. Consultez Gérer les aperçus Databricks.
  2. Interrogez un service de modèle fourni par Databricks à l'aide de ai_query:
SQL
SELECT ai_query(
'databricks-gpt-5-4',
'Summarize the following text: ' || text_column
) AS summary
FROM my_table
LIMIT 10

Les requêtes effectuées via ai_query vers les services de modèles fournis par Databricks sont capturées dans la table système de suivi de l'utilisation (system.ai_gateway.usage). Ces requêtes apparaissent également dans le tableau de bord d’utilisation intégré.

Pour la syntaxe complète ai_query et la référence des paramètres, consultez la fonctionai_query. Pour les bonnes pratiques et les modèles pris en charge, consultez Utiliser ai_query.

Interroger les services de modèle avec des APIs unifiées

Les APIs unifiées offrent une interface compatible avec OpenAI pour interroger les modèles sur Databricks. Utilisez des APIs unifiées pour passer de manière transparente d'un modèle à l'autre entre différents fournisseurs sans modifier votre code.

API MLflow Chat Completions

API MLflow Chat Completions

Python
from openai import OpenAI
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = OpenAI(
api_key=DATABRICKS_TOKEN,
base_url="https://<workspace-url>/ai-gateway/mlflow/v1"
)

chat_completion = client.chat.completions.create(
messages=[
{"role": "user", "content": "Hello!"},
{"role": "assistant", "content": "Hello! How can I assist you today?"},
{"role": "user", "content": "What is Databricks?"},
],
model="<model-service>",
max_tokens=256
)

print(chat_completion.choices[0].message.content)

Remplacez <workspace-url> par l'URL de votre workspace Databricks et <model-service> par le nom complet de votre service de modèle.

API d'embeddings MLflow

API d’embeddings MLflow

Python
from openai import OpenAI
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = OpenAI(
api_key=DATABRICKS_TOKEN,
base_url="https://<workspace-url>/ai-gateway/mlflow/v1"
)

embeddings = client.embeddings.create(
input="What is Databricks?",
model="<model-service>"
)

print(embeddings.data[0].embedding)

Remplacez <workspace-url> par l'URL de votre workspace Databricks et <model-service> par le nom complet de votre service de modèle.

API Supervisor

API Supervisor

La Supervisor API (/mlflow/v1/responses) est une API compatible avec OpenResponseset agnostique du fournisseur pour créer des agents en Bêta. Les administrateurs de compte peuvent activer l'accès à partir de la page Aperçus . Consultez Gérer les aperçus Databricks. Choisissez le meilleur modèle pour votre cas d'utilisation d'agent parmi les fournisseurs, sans modifier votre code.

Python
from openai import OpenAI
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = OpenAI(
api_key=DATABRICKS_TOKEN,
base_url="https://<workspace-url>/ai-gateway/mlflow/v1"
)

response = client.responses.create(
model="<model-service>",
input=[{"role": "user", "content": "What is Databricks?"}]
)

print(response.output_text)

Remplacez <workspace-url> par l'URL de votre workspace Databricks et <model-service> par le nom complet de votre service de modèle.

Interroger les services de modèle avec des APIs natives

Les APIs natives offrent des interfaces spécifiques aux fournisseurs pour query les modèles sur Databricks. Utiliser les APIs natives pour accéder aux dernières fonctionnalités spécifiques au fournisseur.

Chaque API native fonctionne uniquement avec les services de modèles dont le modèle sous-jacent utilise le format d'API correspondant :

  • Utilisez l'OpenAI Responses API pour interroger les services de modèles soutenus par les modèles OpenAI (GPT).
  • Utilisez l'Anthropic Messages API pour interroger les services de modèle pris en charge par les modèles Claude.
  • Utilisez l'Google Gemini API pour interroger les services de modèle reposant sur les modèles Gemini.

Pour interroger un service de modèle quel que soit son modèle sous-jacent, utilisez plutôt les APIs unifiées.

OpenAI Responses API

OpenAI Responses API

Python
from openai import OpenAI
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = OpenAI(
api_key=DATABRICKS_TOKEN,
base_url="https://<workspace-url>/ai-gateway/openai/v1"
)

response = client.responses.create(
model="<model-service>",
max_output_tokens=256,
input=[
{
"role": "user",
"content": [{"type": "input_text", "text": "Hello!"}]
},
{
"role": "assistant",
"content": [{"type": "output_text", "text": "Hello! How can I assist you today?"}]
},
{
"role": "user",
"content": [{"type": "input_text", "text": "What is Databricks?"}]
}
]
)

print(response.output)

Remplacez <workspace-url> par l'URL de votre workspace Databricks et <model-service> par le nom complet de votre service de modèle.

Anthropic Messages API

Anthropic Messages API

Python
import anthropic
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = anthropic.Anthropic(
api_key="unused",
base_url="https://<workspace-url>/ai-gateway/anthropic",
default_headers={
&quot;Authorization&quot;: f&quot;Bearer {DATABRICKS_TOKEN}&quot;,
},
)

message = client.messages.create(
model="<model-service>",
max_tokens=256,
messages=[
{"role": "user", "content": "Hello!"},
{"role": "assistant", "content": "Hello! How can I assist you today?"},
{"role": "user", "content": "What is Databricks?"},
],
)

print(message.content[0].text)

Remplacez <workspace-url> par l'URL de votre workspace Databricks et <model-service> par le nom complet de votre service de modèle.

Google Gemini API

Google Gemini API

Python
from google import genai
from google.genai import types
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = genai.Client(
api_key="databricks",
http_options=types.HttpOptions(
base_url="https://<workspace-url>/ai-gateway/gemini",
headers={
&quot;Authorization&quot;: f&quot;Bearer {DATABRICKS_TOKEN}&quot;,
},
),
)

response = client.models.generate_content(
model="<model-service>",
contents=[
types.Content(
role="user",
parts=[types.Part(text="Hello!")],
),
types.Content(
role="model",
parts=[types.Part(text="Hello! How can I assist you today?")],
),
types.Content(
role="user",
parts=[types.Part(text="What is Databricks?")],
),
],
config=types.GenerateContentConfig(
max_output_tokens=256,
),
)

print(response.text)

Remplacez <workspace-url> par l'URL de votre workspace Databricks et <model-service> par le nom complet de votre service de modèle.

Demandes de tag pour le suivi de l'utilisation

Vous pouvez attacher des tags personnalisés clé-valeur à des requêtes individuelles en utilisant l'en-tête HTTP Databricks-Ai-Gateway-Request-Tags. Les tags de requête sont enregistrés dans la colonne request_tags à la fois dans la table système de suivi de l'utilisation et dans les tables d'inférence, ce qui vous permet de suivre les coûts, d'attribuer l'utilisation et de filtrer l'analytique par projet, équipe, environnement ou toute autre dimension.

La valeur de l’en-tête doit être un objet JSON associant des clés de type chaîne à des valeurs de type chaîne. Par exemple :

JSON
{ "project": "chatbot", "team": "ml-platform", "environment": "production" }

Utilisez le paramètre extra_headers (Python) ou transmettez directement l'en-tête (API REST) pour associer des balises à une requête :

Python
from openai import OpenAI
import json
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = OpenAI(
api_key=DATABRICKS_TOKEN,
base_url="https://<workspace-url>/ai-gateway/mlflow/v1"
)

request_tags = {"project": "chatbot", "team": "ml-platform"}

chat_completion = client.chat.completions.create(
messages=[
{"role": "user", "content": "What is Databricks?"},
],
model="<model-service>",
max_tokens=256,
extra_headers={
&quot;Databricks-Ai-Gateway-Request-Tags&quot;: json.dumps(request_tags)
}
)

Remplacez <workspace-url> par l'URL de votre workspace Databricks et <model-service> par le nom complet de votre service de modèle.

Étapes suivantes