Melakukan kueri pada agen yang disebarkan di Azure Databricks

Pelajari cara mengirim permintaan ke agen yang diterapkan ke Databricks Apps atau endpoint Model Serving. Databricks menyediakan beberapa metode kueri agar sesuai dengan berbagai kasus penggunaan dan kebutuhan integrasi.

Pilih pendekatan kueri yang paling sesuai dengan kasus penggunaan Anda:

Metode	Manfaat utama
Klien Databricks OpenAI (Disarankan)	Integrasi asli, dukungan fitur lengkap, kemampuan streaming
REST API	OpenAI-kompatibel, bebas bahasa, bekerja dengan alat yang ada
Fungsi AI: ai_query	Agen warisan kueri yang kompatibel dengan OpenAI yang dihosting di titik akhir Model Serving saja

Databricks merekomendasikan Klien Databricks OpenAI untuk aplikasi baru. Pilih REST API saat mengintegrasikan dengan platform yang mengharapkan titik akhir yang kompatibel dengan OpenAI.

Klien Databricks OpenAI (Disarankan)

Databricks merekomendasikan agar Anda menggunakan Klien DatabricksOpenAI untuk mengkueri agen yang disebarkan. Bergantung pada API agen yang Anda sebarkan, Anda akan menggunakan respons atau klien penyelesaian obrolan:

Agen yang disebarkan ke Aplikasi

Gunakan contoh berikut untuk agen yang dihosting di Aplikasi Databricks berdasarkan ResponsesAgent antarmuka, yang merupakan pendekatan yang direkomendasikan untuk membangun agen. Anda harus menggunakan token Databricks OAuth untuk mengkueri agen yang dihosting di Aplikasi Databricks.

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)

Jika Anda ingin meneruskan custom_inputs, Anda dapat menambahkannya dengan parameter extra_body.

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)

Untuk mengambil ID pelacakan dari respons, sertakan x-mlflow-return-trace-id header menggunakan extra_headers. Kemudian gunakan MLflow get_trace untuk mengambil jejak lengkap.

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)

Agen pada Penyajian Model

Gunakan contoh berikut untuk agen lama yang dihosting di Model Serving mengikuti ResponsesAgent antarmuka. Anda dapat menggunakan token Databricks OAuth atau Token Akses Pribadi (PAT) untuk mengkueri agen yang dihosting di Model Serving.

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)

Jika Anda ingin meneruskan custom_inputs atau databricks_options, Anda dapat menambahkannya dengan parameter extra_body.

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)

Gunakan contoh berikut untuk agen lama pada penyajian model dengan mengikuti antarmuka ChatAgent atau ChatModel.

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)

Jika Anda ingin meneruskan custom_inputs atau databricks_options, Anda dapat menambahkannya dengan parameter extra_body.

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 menyediakan titik akhir untuk model yang kompatibel dengan OpenAI. Ini memungkinkan Anda menggunakan agen Databricks untuk melayani aplikasi yang memerlukan antarmuka OpenAI.

Pendekatan ini sangat ideal untuk:

• Aplikasi agnostik bahasa yang menggunakan permintaan HTTP
• Mengintegrasikan dengan platform pihak ketiga yang mengharapkan API yang kompatibel dengan OpenAI
• Migrasi dari OpenAI ke Databricks dengan perubahan kode minimal

Autentikasi dengan REST API menggunakan token Databricks OAuth. Lihat Dokumentasi Autentikasi Databricks untuk opsi dan informasi selengkapnya.

Agen yang disebarkan ke Aplikasi

Gunakan contoh berikut untuk agen yang dihosting di Aplikasi Databricks berdasarkan ResponsesAgent antarmuka, yang merupakan pendekatan yang direkomendasikan untuk membangun agen. Anda harus menggunakan token Databricks OAuth untuk mengkueri agen yang dihosting di Aplikasi Databricks.

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

Jika Anda ingin meneruskan custom_inputs, Anda dapat menambahkannya ke isi permintaan:

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

Untuk mengambil ID pelacakan dari respons, sertakan x-mlflow-return-trace-id header dalam permintaan Anda. Isi respons mencakup bidang metadata.trace_id yang berisi ID pelacakan. Untuk permintaan streaming, ID jejak dikirim sebagai peristiwa SSE terpisah (data: {"trace_id": "tr-..."}) di dekat akhir stream. Kemudian gunakan MLflow get_trace untuk mengambil jejak lengkap menggunakan ID pelacakan.

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

Agen pada Penyajian Model

Gunakan contoh berikut untuk agen lama yang dihosting di Model Serving mengikuti ResponsesAgent antarmuka. Anda dapat menggunakan token Databricks OAuth atau Token Akses Pribadi (PAT) untuk mengkueri agen yang dihosting di Model Serving. Panggilan REST API setara dengan:

• Menggunakan Klien Databricks OpenAI dengan responses.create.
• Mengirim permintaan POST ke URL titik akhir tertentu (misalnya: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Untuk informasi selengkapnya, lihat halaman Model Serving endpoint Anda dan Dokumentasi Model Serving.

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

Jika Anda ingin meneruskan custom_inputs atau databricks_options, Anda dapat menambahkannya ke isi permintaan:

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

Gunakan yang berikut ini untuk agen yang dibuat dengan antarmuka ChatAgent atau ChatModel lama. Ini setara dengan:

• Menggunakan Klien Databricks OpenAI dengan chat.completions.create.
• Mengirim permintaan POST ke URL titik akhir tertentu (misalnya: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Untuk informasi selengkapnya, lihat halaman Model Serving endpoint Anda dan Dokumentasi Model Serving.

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

Jika Anda ingin meneruskan custom_inputs atau databricks_options, Anda dapat menambahkannya ke isi permintaan:

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

Fungsi AI: ai_query

Anda dapat menggunakan ai_query untuk mengkueri agen yang disebarkan yang dihosting pada model yang melayani menggunakan SQL. Lihat ai_query fungsi untuk sintaksis SQL dan definisi parameter.

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

Langkah berikutnya

Memantau aplikasi GenAI dalam lingkungan produksi