Fazer query de um agente implantado no Databricks

Saiba como enviar solicitações para agentes implantados em Databricks Apps ou em Endpoint de Model Serving. O Databricks oferece vários métodos de query para atender a diferentes casos de uso e necessidades de integração.

Selecione a abordagem de query que melhor se adapta ao seu caso de uso:

Método	Principais benefícios
Cliente Databricks OpenAI (Recomendado)	Integração nativa, suporte completo a recursos, funcionalidades de transmissão
API REST	Compatível com OpenAI, independente de linguagem, funciona com ferramentas existentes.
AI Functions: ai_query	Compatível com OpenAI, consultar apenas agentes legados hospedados em endpoints de Model Serving

Método	Principais benefícios
Cliente Databricks OpenAI (Recomendado)	Integração nativa, suporte completo a recursos, funcionalidades de transmissão
API REST	Compatível com OpenAI, independente de linguagem, funciona com ferramentas existentes.
AI Functions: ai_query	Compatível com OpenAI, consultar apenas agentes legados hospedados em endpoints de Model Serving

A Databricks recomenda o **Databricks OpenAI Client** para novas aplicações. Escolha a REST API ao integrar com plataformas que esperam OpenAI-compatíveis Endpoint.

Cliente OpenAI da Databricks (Recomendado)

O Databricks recomenda o uso do Cliente DatabricksOpenAI para query um agente implantado. Dependendo da API do seu agente implantado, serão usadas as respostas ou o cliente de conclusões de chat:

Agents deployed to Apps

Use o exemplo a seguir para agentes hospedados no Databricks Apps seguindo a interface ResponsesAgent, que é a abordagem recomendada para criar agentes. Você deve usar um token OAuth do Databricks para fazer query de agentes hospedados no Databricks Apps.

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)

Se você quiser passar custom_inputs, você pode adicioná-los com o parâmetro 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)

Para recuperar um ID de rastreamento da resposta, inclua o cabeçalho x-mlflow-return-trace-id usando extra_headers. Em seguida, use o MLflow get_trace para recuperar o rastreamento completo.

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

Use o exemplo a seguir para agentes legados hospedados no Model Serving seguindo a interface ResponsesAgent. Você pode usar um token OAuth do Databricks ou um Personal Access Token (PAT) para consultar agentes hospedados no Model Serving.

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)

Se quiser passar custom_inputs ou databricks_options, pode adicioná-los com o parâmetro 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)

Use o seguinte exemplo para agentes legados em servindo modelo, seguindo as interfaces ChatAgent ou 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)

Se quiser passar custom_inputs ou databricks_options, pode adicioná-los com o parâmetro 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)

API REST

A API REST do Databricks fornece Endpoint para modelos que são compatíveis com OpenAI. Isso permite usar agentes do Databricks para servir aplicativos que exigem interfaces OpenAI.

Esta abordagem é ideal para:

• Aplicativos independentes de linguagem que usam solicitações HTTP
• Integrando com plataformas de terceiros que esperam APIs compatíveis com OpenAI
• Migrando da OpenAI para a Databricks com mínimas alterações de código.

Autenticar-se com a API REST usando um token OAuth do Databricks. Consulte a Documentação de Autenticação do Databricks para mais opções e informações.

Agents deployed to Apps

Use o exemplo a seguir para agentes hospedados no Databricks Apps seguindo a interface ResponsesAgent, que é a abordagem recomendada para criar agentes. Você deve usar um token OAuth do Databricks para fazer query de agentes hospedados no Databricks Apps.

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

Se o senhor quiser passar custom_inputs, o senhor poderá adicioná-los ao corpo da solicitação:

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

Para recuperar um ID de rastreamento da resposta, inclua o cabeçalho x-mlflow-return-trace-id em sua solicitação. O corpo da resposta inclui um campo metadata.trace_id contendo o ID do rastreamento. Para solicitações de transmissão, o ID do rastreamento é enviado como um evento SSE separado (data: {"trace_id": "tr-..."}) perto do final da transmissão. Em seguida, use o MLflow get_trace para recuperar o rastreamento completo usando o ID do rastreamento.

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

Use o exemplo a seguir para agentes legados hospedados no Model Serving seguindo a interface ResponsesAgent. O senhor pode usar tanto um token OAuth da Databricks quanto um Token de Acesso Pessoal (PAT) para consultar agentes hospedados no Model Serving. A chamada de API REST é equivalente a:

• Usando o cliente Databricks OpenAI com responses.create.
• Enviando uma solicitação POST para o URL do endpoint específico (ex: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Para mais informações, consulte a página Model Serving do seu endpoint e a Documentação do 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
  }'

Se você quiser passar custom_inputs ou databricks_options, pode adicioná-los ao corpo da solicitação:

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

Use o seguinte para agentes criados com interfaces legadas do ChatAgent ou ChatModel. Isso é equivalente a:

• Usando o cliente Databricks OpenAI com chat.completions.create.
• Enviando uma solicitação POST para o URL do endpoint específico (ex: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Para mais informações, consulte a página Model Serving do seu endpoint e a Documentação do 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
  }'

Se você quiser passar custom_inputs ou databricks_options, pode adicioná-los ao corpo da solicitação:

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

Você pode usar ai_query para query um agente implantado hospedado no servindo modelo usando SQL. Consulte a funçãoai_query para a sintaxe SQL e definições de parâmetros.

SQL

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

Próximos passos

Monitore aplicativos GenAI em produção