Servicios del modelo de consulta

Importante

Esta característica se encuentra en su versión beta. Los administradores de cuentas pueden controlar el acceso a esta característica desde la página Vista previa de la consola de la cuenta. Consulte Administrar versiones preliminares de Azure Databricks.

En esta página se describe cómo consultar los servicios de modelo en el catálogo de Unity mediante las API admitidas.

Nota:

Cuando un servicio de modelo dirige a un destino de servicio del proveedor del modelo, solo se aplican las funciones de Unity AI Gateway del servicio de modelo (como límites de velocidad, barreras de seguridad, tablas de inferencia y alternativas). Se omite cualquier función de Unity AI Gateway configurada en el propio servicio del proveedor de modelos.

Requirements

APIs e integraciones admitidas

Unity AI Gateway admite las siguientes API e integraciones:

Consultar los servicios del modelo con ai_query

Puede usar la ai_query función para consultar los servicios de modelo proporcionados por Azure Databricks directamente desde SQL o Python. Esto le permite capturar información de seguimiento del uso para las cargas de trabajo de inferencia por lotes.

Nota:

  • ai_queryLa compatibilidad con Unity AI Gateway solo está disponible para los servicios de modelo proporcionados por Azure Databricks (por ejemplo, databricks-gpt-5-4 o databricks-claude-sonnet-4). Los servicios de modelo que cree en Unity AI Gateway aún no se admiten.
  • Solo el seguimiento del uso se aplica a ai_query las cargas de trabajo de inferencia por lotes. Otras características de Unity AI Gateway, como los límites de tasa, las salvaguardas, las tablas de inferencia y los mecanismos de respaldo, no se aplican.

Primeros pasos:

  1. Habilite la versión preliminar de Unity AI Gateway para su cuenta. Consulte Administrar versiones preliminares de Azure Databricks.
  2. Consulte un servicio de modelo proporcionado por Azure Databricks mediante ai_query:
SELECT ai_query(
  'databricks-gpt-5-4',
  'Summarize the following text: ' || text_column
) AS summary
FROM my_table
LIMIT 10

Las solicitudes realizadas a través ai_query de los servicios de modelo proporcionados por Azure Databricks se capturan en la tabla del sistema de seguimiento de uso (system.ai_gateway.usage). Estas solicitudes también aparecen en el panel de uso integrado.

Para obtener una sintaxis completa ai_query y una referencia de parámetros, consulte ai_query function. Para conocer los procedimientos recomendados y los modelos admitidos, consulte Uso de ai_query.

Consulta de servicios de modelo con API unificadas

Las API unificadas ofrecen una interfaz compatible con OpenAI para consultar modelos en Azure Databricks. Use API unificadas para cambiar sin problemas entre modelos de diferentes proveedores sin cambiar el código.

API de finalizaciones de chat de MLflow

API de finalizaciones de chat de 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"
)

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)

REST API

curl \
  -u token:$DATABRICKS_TOKEN \
  -X POST \
  -H "Content-Type: application/json" \
  -d '{
    "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?"}
    ]
  }' \
  https://<workspace-url>/ai-gateway/mlflow/v1/chat/completions

Reemplace por <workspace-url> la dirección URL del área de trabajo de Azure Databricks y <model-service> por el nombre completo del servicio de modelo.

API de inserción de MLflow

API de inserción de 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)

REST API

curl \
  -u token:$DATABRICKS_TOKEN \
  -X POST \
  -H "Content-Type: application/json" \
  -d '{
    "model": "<model-service>",
    "input": "What is Databricks?"
  }' \
  https://<workspace-url>/ai-gateway/mlflow/v1/embeddings

Reemplace por <workspace-url> la dirección URL del área de trabajo de Azure Databricks y <model-service> por el nombre completo del servicio de modelo.

Supervisor API

Supervisor API

La API supervisor (/mlflow/v1/responses) es una API independiente del proveedor y compatible con OpenResponses para compilar agentes en Beta. Los administradores de cuentas pueden habilitar el acceso desde la página Vistas previas . Consulte Administrar versiones preliminares de Azure Databricks. Elija el mejor modelo para el caso de uso del agente entre proveedores, sin cambiar el código.

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)

REST API

curl \
  -u token:$DATABRICKS_TOKEN \
  -X POST \
  -H "Content-Type: application/json" \
  -d '{
    "model": "<model-service>",
    "input": [
      {"role": "user", "content": "What is Databricks?"}
    ]
  }' \
  https://<workspace-url>/ai-gateway/mlflow/v1/responses

Reemplace por <workspace-url> la dirección URL del área de trabajo de Azure Databricks y <model-service> por el nombre completo del servicio de modelo.

Consulta de servicios de modelo con API nativas

Las API nativas ofrecen interfaces específicas del proveedor para consultar modelos en Azure Databricks. Use las API nativas para acceder a las características específicas del proveedor más recientes.

Cada API nativa solo funciona con servicios de modelo cuyo modelo subyacente usa el formato de API coincidente:

Para consultar un servicio de modelo independientemente de su modelo subyacente, use las API unificadas en su lugar.

API de respuestas de OpenAI

API de respuestas de OpenAI

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)

REST API

curl \
  -u token:$DATABRICKS_TOKEN \
  -X POST \
  -H "Content-Type: application/json" \
  -d '{
    "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?"}]
      }
    ]
  }' \
  https://<workspace-url>/ai-gateway/openai/v1/responses

Reemplace por <workspace-url> la dirección URL del área de trabajo de Azure Databricks y <model-service> por el nombre completo del servicio de modelo.

API de mensajes de Anthropic

API de mensajes de Anthropic

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={
    "Authorization": f"Bearer {DATABRICKS_TOKEN}",
  },
)

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)

REST API

curl \
  -u token:$DATABRICKS_TOKEN \
  -X POST \
  -H "Content-Type: application/json" \
  -d '{
    "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?"}
    ]
  }' \
  https://<workspace-url>/ai-gateway/anthropic/v1/messages

Reemplace por <workspace-url> la dirección URL del área de trabajo de Azure Databricks y <model-service> por el nombre completo del servicio de modelo.

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={
      "Authorization": f"Bearer {DATABRICKS_TOKEN}",
    },
  ),
)

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)

REST API

curl \
  -u token:$DATABRICKS_TOKEN \
  -X POST \
  -H "Content-Type: application/json" \
  -d '{
    "contents": [
      {
        "role": "user",
        "parts": [{"text": "Hello!"}]
      },
      {
        "role": "model",
        "parts": [{"text": "Hello! How can I assist you today?"}]
      },
      {
        "role": "user",
        "parts": [{"text": "What is Databricks?"}]
      }
    ],
    "generationConfig": {
      "maxOutputTokens": 256
    }
  }' \
  https://<workspace-url>/ai-gateway/gemini/v1beta/models/<model-service>:generateContent

Reemplace por <workspace-url> la dirección URL del área de trabajo de Azure Databricks y <model-service> por el nombre completo del servicio de modelo.

Solicitudes de etiquetas para el seguimiento de uso

Puede adjuntar etiquetas de clave-valor personalizadas a solicitudes individuales mediante el Databricks-Ai-Gateway-Request-Tags encabezado HTTP. Las etiquetas de solicitud se registran en la request_tags columna tanto en la tabla del sistema de seguimiento de uso como en las tablas de inferencia, lo que permite realizar un seguimiento de los costos, el uso de atributos y el análisis de filtros por proyecto, equipo, entorno o cualquier otra dimensión.

El valor del encabezado debe ser un objeto JSON que asigna claves de cadena a valores de cadena. Por ejemplo:

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

Use el parámetro extra_headers (Python) o pase el encabezado directamente (API REST) para adjuntar etiquetas a una solicitud:

Python (SDK de OpenAI)

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={
    "Databricks-Ai-Gateway-Request-Tags": json.dumps(request_tags)
  }
)

Python (SDK de Anthropic)

import anthropic
import json
import os

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

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

client = anthropic.Anthropic(
  api_key="unused",
  base_url="https://<workspace-url>/ai-gateway/anthropic",
  default_headers={
    "Authorization": f"Bearer {DATABRICKS_TOKEN}",
    "Databricks-Ai-Gateway-Request-Tags": json.dumps(request_tags),
  },
)

message = client.messages.create(
  model="<model-service>",
  max_tokens=256,
  messages=[
    {"role": "user", "content": "What is Databricks?"},
  ],
)

REST API

curl \
  -u token:$DATABRICKS_TOKEN \
  -X POST \
  -H "Content-Type: application/json" \
  -H 'Databricks-Ai-Gateway-Request-Tags: {"project": "chatbot", "team": "ml-platform"}' \
  -d '{
    "model": "<model-service>",
    "max_tokens": 256,
    "messages": [
      {"role": "user", "content": "What is Databricks?"}
    ]
  }' \
  https://<workspace-url>/ai-gateway/mlflow/v1/chat/completions

Reemplace por <workspace-url> la dirección URL del área de trabajo de Azure Databricks y <model-service> por el nombre completo del servicio de modelo.

Pasos siguientes