Запрос агента, развернутого в Azure Databricks

Узнайте, как отправлять запросы агентам, развернутыми в приложениях Databricks или конечных точках обслуживания моделей. Databricks предоставляет несколько методов запросов в соответствии с различными вариантами использования и потребностями интеграции.

Выберите подход к запросу, который лучше всего подходит для вашего варианта использования:

Метод Основные преимущества
Клиент Databricks OpenAI (рекомендуется) Встроенная интеграция, полная поддержка функций, возможности потоковой передачи
REST API Совместимый с OpenAI, независимый от языка, работает с существующими инструментами.
Функции ИИ: ai_query Совместимый с OpenAI, запрос устаревших агентов, размещенных только в конечных точках обслуживания моделей

Databricks рекомендует клиент Databricks OpenAI для новых приложений. Выберите REST API при интеграции с платформами, которые ожидают конечных точек, совместимых с OpenAI.

Databricks рекомендует использовать клиент DatabricksOpenAI для запроса развернутого агента. В зависимости от API развернутого агента вы будете использовать ответы или клиент для завершения диалога:

Агенты, развернутые в приложениях

Используйте следующий пример для агентов, размещенных в Databricks Apps , следуя ResponsesAgent интерфейсу, который является рекомендуемым подходом для создания агентов. Для запроса агентов, размещенных в приложениях Databricks, необходимо использовать токен OAuth 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)

Если вы хотите передать custom_inputs, вы можете добавить их с помощью параметра 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)

Чтобы получить идентификатор трассировки из ответа, включите x-mlflow-return-trace-id заголовок с помощью extra_headers. Затем используйте MLflow get_trace для получения полной трассировки.

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)

Агенты в службе моделей

Используйте следующий пример для устаревших агентов, размещенных в службе модели после ResponsesAgent интерфейса. Вы можете использовать токен OAuth Databricks или Персональный токен доступа (PAT) для выполнения запросов к агентам, размещенным в сервисе модели.

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.

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.

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.

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

REST API Databricks предоставляет конечные точки для моделей, совместимых с OpenAI. Это позволяет использовать агенты Databricks для обслуживания приложений, требующих интерфейсов OpenAI.

Этот подход идеально подходит для:

  • Приложения, не зависящие от языка, использующие HTTP-запросы
  • Интеграция с сторонними платформами, которые ожидают API, совместимые с OpenAI
  • Миграция из OpenAI в Databricks с минимальными изменениями кода

Аутентификация через REST API с использованием токена OAuth Databricks. Для получения дополнительных сведений и вариантов обратитесь к документации по аутентификации Databricks.

Агенты, развернутые в приложениях

Используйте следующий пример для агентов, размещенных в Databricks Apps , следуя ResponsesAgent интерфейсу, который является рекомендуемым подходом для создания агентов. Для запроса агентов, размещенных в приложениях Databricks, необходимо использовать токен OAuth 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
  }'

Если вы хотите передать custom_inputs, вы можете добавить его в тело запроса:

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

Чтобы получить идентификатор трассировки из ответа, добавьте заголовок x-mlflow-return-trace-id в этот запрос. Текст ответа содержит поле metadata.trace_id, содержащее идентификатор трассировки. Для потоковых запросов идентификатор трассировки отправляется как отдельное событие SSE (data: {"trace_id": "tr-..."}) в конце потока. Затем используйте MLflow get_trace для получения полной трассировки с помощью идентификатора трассировки.

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

Агенты в службе моделей

Используйте следующий пример для устаревших агентов, размещенных в службе модели после ResponsesAgent интерфейса. Вы можете использовать токен OAuth Databricks или Персональный токен доступа (PAT) для выполнения запросов к агентам, размещенным в сервисе модели. Вызов REST API эквивалентен следующим:

  • Использование клиента Databricks OpenAI с responses.create.
  • Отправка запроса POST на URL-адрес конкретной конечной точки (например: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Дополнительные сведения см. на странице обслуживания модели конечной точки и документации по обслуживанию моделей.
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, вы можете добавить их в текст запроса:

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. Это эквивалентно следующему:

  • Использование клиента Databricks OpenAI с chat.completions.create.
  • Отправка запроса POST на URL-адрес конкретной конечной точки (например: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Дополнительные сведения см. на странице обслуживания модели конечной точки и документации по обслуживанию моделей.
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, вы можете добавить их в текст запроса:

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_query

Вы можете запросить ai_query развернутый агент, размещенный в службе моделей с помощью SQL. См. ai_query функцию для определений синтаксиса и параметров SQL.

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

Дальнейшие действия

Мониторинг приложений GenAI в рабочей среде