Databricks にデプロイされたエージェントをクエリーする

Databricks Apps または Model Serving エンドポイントにデプロイされたエージェントにリクエストを送信する方法について説明します。Databricksは、さまざまなユースケースと統合ニーズに対応するために複数のクエリメソッドを提供しています。

ユースケースに最適なクエリ手法を選択してください：

手法	主なメリット
Databricks OpenAI クライアント (推奨)	ネイティブ統合、フル機能サポート、ストリーミング機能
REST API	OpenAI対応、言語に依存せず、既存のツールで動作します。
AI Functions: ai_query	OpenAI と互換性があり、Model Serving エンドポイントでホストされているレガシーエージェントのみをクエリします。

Databricksは、新規アプリケーションに**Databricks OpenAI Client**を推奨しています。 OpenAI互換のエンドポイントを想定しているプラットフォームと統合する際には、**REST API** を選択してください。

Databricks OpenAI クライアント（推奨）

Databricksでは、デプロイされたエージェントをクエリするためにDatabricks OpenAI Clientを使用することをお勧めします。デプロイされたエージェントのAPIに応じて、レスポンスまたはチャット補完クライアントのいずれかを使用します。

Agents deployed to Apps

エージェントを構築するための推奨アプローチであるResponsesAgentインターフェースに従った、Databricks Appsでホストされているエージェントには、次の例を使用してください。Databricks Appsでホストされているエージェントをクエリするには、Databricks OAuth トークンを使用する必要があります。

Python

from databricks.sdk import WorkspaceClient
from databricks_openai import DatabricksOpenAI

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

# The WorkspaceClient must be configured with OAuth authentication
# See: https://docs.databricks.com/aws/en/dev-tools/auth/oauth-u2m.html
w = WorkspaceClient()

client = DatabricksOpenAI(workspace_client=w)

# Run for non-streaming responses. Calls the "invoke" method
# Include the "apps/" prefix in the model name
response = client.responses.create(model=f"apps/{app_name}", input=input_msgs)
print(response)

# Include stream=True for streaming responses. Calls the "stream" method
# Include the "apps/" prefix in the model name
streaming_response = client.responses.create(
    model=f"apps/{app_name}", input=input_msgs, stream=True
)
for chunk in streaming_response:
    print(chunk)

custom_inputsを渡したい場合は、extra_bodyパラメーターで追加できます。

Python

streaming_response = client.responses.create(
    model=f"apps/{app_name}",
    input=input_msgs,
    stream=True,
    extra_body={
        "custom_inputs": {"id": 5},
    },
)
for chunk in streaming_response:
    print(chunk)

応答からトレースIDを取得するには、extra_headersを使用してx-mlflow-return-trace-idヘッダーを含めます。次に、MLflow get_trace を使用して、完全なトレースを取得します。

Python

response = client.responses.create(
    model=f"apps/{app_name}",
    input=input_msgs,
    extra_headers={"x-mlflow-return-trace-id": "true"},
)
trace_id = response.metadata["trace_id"]
trace = client.get_trace(trace_id)

Agents on Model Serving

ResponsesAgent インターフェースに従うレガシーな Model Serving にホストされているエージェントに対して、次の例を使用してください。Model Servingでホストされているエージェントをクエリするには、Databricks OAuthトークンまたは個人用アクセストークン (PAT) のいずれかを使用できます。

Python

from databricks_openai import DatabricksOpenAI

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

client = DatabricksOpenAI()

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

custom_inputs または databricks_options を渡す場合は、extra_body パラメーターで追加できます：

Python

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)

ChatAgent または ChatModel インターフェースに従うモデルサービングのレガシーエージェントには、以下の例を使用してください。

Python

from databricks.sdk import WorkspaceClient

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

ws_client = WorkspaceClient()
client = ws_client.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)

custom_inputs または databricks_options を渡す場合は、extra_body パラメーターで追加できます：

Python

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)

REST API

Databricks REST API は、OpenAI と互換性のあるモデルのエンドポイントを提供します。これにより、Databricks エージェントを使用して、OpenAI インターフェースを必要とするアプリケーションを提供できます。

このアプローチは、以下に最適です：

• HTTPリクエストを使用する言語非依存のアプリケーション
• OpenAI互換のAPIsを期待するサードパーティプラットフォームとの連携
• 最小限のコード変更でOpenAIからDatabricksへの移行

Databricks OAuth トークンを使用して REST API に認証します。その他のオプションと情報については、Databricks 認証ドキュメントを参照してください。

Agents deployed to Apps

エージェントを構築するための推奨アプローチであるResponsesAgentインターフェースに従った、Databricks Appsでホストされているエージェントには、次の例を使用してください。Databricks Appsでホストされているエージェントをクエリするには、Databricks OAuth トークンを使用する必要があります。

Bash

curl --request POST \
  --url <app-url>.databricksapps.com/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "input": [{ "role": "user", "content": "hi" }],
    "stream": true
  }'

custom_inputs を渡したい場合は、リクエストボディに追加できます。

Bash

curl --request POST \
  --url <app-url>.databricksapps.com/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "input": [{ "role": "user", "content": "hi" }],
    "stream": true,
    "custom_inputs": { "id": 5 }
  }'

応答からトレースIDを取得するには、リクエストにx-mlflow-return-trace-idヘッダーを含めます。応答本文には、トレースIDを含むmetadata.trace_idフィールドが含まれます。ストリーミングリクエストの場合、トレースIDはストリームの終了近くで個別のSSEイベント(data: {"trace_id": "tr-..."})として送信されます。次に、トレースIDを使用してMLflow get_traceで完全なトレースを取得します。

Bash

curl --request POST \
  --url <app-url>.databricksapps.com/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --header 'x-mlflow-return-trace-id: true' \
  --data '{
    "input": [{ "role": "user", "content": "hi" }]
  }'

Agents on Model Serving

レガシーのモデルサービングでホストされているエージェントには、ResponsesAgentインターフェイスに従って次の例を使用してください。Databricks OAuthトークンまたは個人アクセストークン(PAT)のいずれかを使用して、モデルサービングでホストされているエージェントをクエリできます。The REST API call is equivalent to:

• responses.createでDatabricks OpenAIクライアントを使用しています。
• 特定のエンポイントのURLにPOSTリクエストを送信します (例: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations)。詳細については、お使いのエンドポイントのModel ServingページおよびModel Servingドキュメントを参照してください。

Bash

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

custom_inputsまたはdatabricks_optionsを渡したい場合は、リクエストボディに追加できます。

Bash

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

レガシーなChatAgentまたはChatModelインターフェースで作成されたエージェントには、以下を使用してください。これは以下と同じです。

• chat.completions.createでDatabricks OpenAIクライアントを使用しています。
• 特定のエンポイントのURLにPOSTリクエストを送信します (例: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations)。詳細については、お使いのエンドポイントのModel ServingページおよびModel Servingドキュメントを参照してください。

Bash

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

custom_inputsまたはdatabricks_optionsを渡したい場合は、リクエストボディに追加できます。

Bash

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

AI Functions: ai_query

ai_queryを使用して、モデルサービングでホストされているデプロイ済みエージェントをSQLでクエリできます。SQL構文とパラメーターの定義については、ai_query関数を参照してください。

SQL

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

次のステップ

本番運用でGenAIアプリを監視します。