Pular para o conteúdo principal

Consultar serviços de modelo

info

Beta

Esse recurso está em Beta. Administradores da conta podem controlar o acesso a este recurso na página Prévias do console da conta. Consulte Gerenciar prévias do Databricks.

Esta página descreve como consultar serviços de modelo no Unity Catalog usando APIs compatíveis.

Requisitos

APIs e Integrações Suportadas

O Unity AI Gateway suporta as seguintes APIs e integrações:

Consultar serviços de modelo com ai_query

É possível usar a ai_query função para consultar os serviços de modelo fornecidos pelo Databricks diretamente do SQL ou Python. Isso permite capturar informações de acompanhamento de uso para suas cargas de trabalho de inferência em lote.

nota
  • ai_query O suporte para o Unity AI Gateway está disponível apenas para serviços de modelo fornecidos pela Databricks (por exemplo, databricks-gpt-5-4 ou databricks-claude-sonnet-4). Os serviços de modelo que você cria no Unity AI Gateway ainda não são compatíveis.
  • Somente o acompanhamento de uso se aplica às cargas de trabalho de inferência em lote ai_query. Outros recursos do Unity AI Gateway, como limites de taxa, guardrails, tabelas de inferência e fallbacks, não se aplicam.

Para começar:

  1. Habilite a pré-visualização do Unity AI Gateway para sua account. Consulte Gerenciar prévias do Databricks.
  2. Consulte um serviço de modelo fornecido pelo Databricks usando ai_query:
SQL
SELECT ai_query(
'databricks-gpt-5-4',
'Summarize the following text: ' || text_column
) AS summary
FROM my_table
LIMIT 10

As solicitações feitas por meio de ai_query para serviços de modelo fornecidos pelo Databricks são capturadas na tabela do sistema de acompanhamento de uso (system.ai_gateway.usage). Essas solicitações também aparecem no painel de uso integrado.

Para a sintaxe completa ai_query e referência de parâmetros, consulte a ai_query função. Para melhores práticas e modelos compatíveis, consulte Use ai_query.

Consultar serviços de modelo com APIs unificadas

As APIs unificadas oferecem uma interface compatível com OpenAI para consultar modelos no Databricks. Use APIs unificadas para alternar facilmente entre modelos de diferentes provedores sem alterar seu código.

API de Conclusões de Chat do MLflow

API de Conclusões de Chat do 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)

Substitua <workspace-url> pelo URL do seu workspace do Databricks e <model-service> pelo nome totalmente qualificado do seu serviço de modelo.

API de Embeddings MLflow

API de Embeddings do 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)

Substitua <workspace-url> pelo URL do seu workspace do Databricks e <model-service> pelo nome totalmente qualificado do seu serviço de modelo.

API do Supervisor

API do Supervisor

Consultar serviços de modelo com APIs nativas

APIs nativas oferecem interfaces específicas do provedor para consultar modelos no Databricks. Use APIs nativas para acessar os recursos específicos mais recentes do provedor.

API de Respostas da OpenAI

API de respostas da 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)

Substitua <workspace-url> pelo URL do seu workspace do Databricks e <model-service> pelo nome totalmente qualificado do seu serviço de modelo.

API de Mensagens da Anthropic

API de Mensagens do 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={
&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)

Substitua <workspace-url> pelo URL do seu workspace do Databricks e <model-service> pelo nome totalmente qualificado do seu serviço de modelo.

API Gemini do Google

API do Google Gemini

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)

Substitua <workspace-url> pelo URL do seu workspace do Databricks e <model-service> pelo nome totalmente qualificado do seu serviço de modelo.

Solicitações de tag para acompanhamento de uso

Você pode anexar tags key-value personalizadas a solicitações individuais usando o cabeçalho HTTP Databricks-Ai-Gateway-Request-Tags. As tags de solicitação são registradas em log na coluna request_tags tanto na tabela do sistema de acompanhamento de uso quanto nas tabelas de inferência, permitindo que você acompanhe custos, atribua uso e filtre analítica por projeto, equipe, ambiente ou qualquer outra dimensão.

O valor do cabeçalho deve ser um objeto JSON que mapeia key de strings para valores de strings. Por exemplo:

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

Use o parâmetro extra_headers (Python) ou passe o cabeçalho diretamente (API REST) para anexar tags a uma solicitação:

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)
}
)

Substitua <workspace-url> pelo URL do seu workspace do Databricks e <model-service> pelo nome totalmente qualificado do seu serviço de modelo.

Próximos os passos