メインコンテンツまでスキップ
に最終更新

AIゲートウェイエンドポイントのクエリ

備考

ベータ版

この機能はベータ版です。アカウント管理者は、アカウント コンソールの [プレビュー] ページからこの機能へのアクセスを制御できます。 「Databricks プレビューの管理」を参照してください。

このページでは、サポートされているAPIsを使用してAI Gateway (ベータ版)エンドポイントをクエリする方法について説明します。

要件​

サポートされているAPIsと統合​

AI Gateway は次のAPIsと統合をサポートしています。

統合APIsでエンドポイントをクエリする​

統合APIs 、 Databricks上のモデルをクエリするためのOpenAI互換インターフェースを提供します。 統合されたAPIsを使用すると、コードを変更せずに、さまざまなプロバイダーのモデル間をシームレスに切り替えることができます。

MLflow チャット完了 API​

Python
from openai import OpenAI
import os

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

client = OpenAI(
  api_key=DATABRICKS_TOKEN,
  base_url="https://<ai-gateway-url>/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="<ai-gateway-endpoint>",
  max_tokens=256
)

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

<ai-gateway-url> AI ゲートウェイ URL に置き換え、 <ai-gateway-endpoint> AI ゲートウェイのエンドポイント名に置き換えます。

MLflow 埋め込み API​

Python
from openai import OpenAI
import os

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

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

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

print(embeddings.data[0].embedding)

<ai-gateway-url> AI ゲートウェイ URL に置き換え、 <ai-gateway-endpoint> AI ゲートウェイのエンドポイント名に置き換えます。

ネイティブAPIsを使用してエンドポイントをクエリする​

ネイティブAPIs Databricks上のモデルをクエリするためのプロバイダー固有のインターフェイスを提供します。 ネイティブAPIsを使用して、プロバイダー固有の最新機能にアクセスします。

OpenAI レスポンス API​

Python
from openai import OpenAI
import os

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

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

response = client.responses.create(
  model="<ai-gateway-endpoint>",
  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)

<ai-gateway-url> AI ゲートウェイ URL に置き換え、 <ai-gateway-endpoint> AI ゲートウェイのエンドポイント名に置き換えます。

AnthropicメッセージAPI​

Python
import anthropic
import os

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

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

message = client.messages.create(
  model="<ai-gateway-endpoint>",
  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)

<ai-gateway-url> AI ゲートウェイ URL に置き換え、 <ai-gateway-endpoint> AI ゲートウェイのエンドポイント名に置き換えます。

ジェミニ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://<ai-gateway-url>/gemini",
    headers={
      "Authorization": f"Bearer {DATABRICKS_TOKEN}",
    },
  ),
)

response = client.models.generate_content(
  model="<ai-gateway-endpoint>",
  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)

<ai-gateway-url> AI ゲートウェイ URL に置き換え、 <ai-gateway-endpoint> AI ゲートウェイのエンドポイント名に置き換えます。

次のステップ​