Consultar um agente implantado no Azure Databricks

Saiba como enviar solicitações para agentes implantados nos Aplicativos do Databricks ou nos endpoints de Serviço de Modelos. O Databricks fornece vários métodos de consulta para atender a diferentes casos de uso e necessidades de integração.

Selecione a abordagem de consulta que melhor se ajusta ao seu caso de uso:

O Databricks recomenda o Cliente Do Databricks OpenAI para novos aplicativos. Escolha a API REST ao integrar com plataformas que esperam pontos de extremidade compatíveis com OpenAI.

Cliente OpenAI do Databricks (recomendado)

O Databricks recomenda que você use o cliente DatabricksOpenAI para consultar um agente implantado. Dependendo da API do agente implantado, você usará as respostas ou o cliente de conclusão de chat:

Agentes implantados nos Aplicativos

Use o exemplo a seguir para agentes hospedados nos Aplicativos do Databricks seguindo a ResponsesAgent interface, que é a abordagem recomendada para a criação de agentes. Você deve usar um token OAuth do Databricks para consultar agentes hospedados nos Aplicativos do 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)

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

Para recuperar uma 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.

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)

Agentes no Serviço de Modelo

Use o exemplo a seguir para agentes herdados hospedados no Serviço de Modelo seguindo a interface ResponsesAgent. Você pode usar um token OAuth do Databricks ou um PAT (Token de Acesso Pessoal) para consultar agentes hospedados no Serviço de Modelo.

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 você quiser passar custom_inputs ou databricks_options, poderá adicioná-los com o extra_body parâmetro:

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 exemplo a seguir para agentes herdados no Serviço de Modelo seguindo as interfaces ChatAgent ou 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)

Se você quiser passar custom_inputs ou databricks_options, poderá adicioná-los com o extra_body parâmetro:

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 REST API do Databricks fornece endpoints para modelos compatíveis com OpenAI. Isso permite que você use agentes do Databricks para atender a aplicativos que exigem interfaces OpenAI.

Essa abordagem é ideal para:

• Aplicativos independentes de linguagem que usam solicitações HTTP
• Integração com plataformas de terceiros que esperam APIs compatíveis com OpenAI
• Migrando do OpenAI para o Databricks com alterações mínimas de código

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

Agentes implantados nos Aplicativos

Use o exemplo a seguir para agentes hospedados nos Aplicativos do Databricks seguindo a ResponsesAgent interface, que é a abordagem recomendada para a criação de agentes. Você deve usar um token OAuth do Databricks para consultar agentes hospedados nos Aplicativos do 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
  }'

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

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 uma ID de rastreamento da resposta, inclua o x-mlflow-return-trace-id cabeçalho em sua solicitação. O corpo da resposta inclui um metadata.trace_id campo que contém a ID de rastreamento. Para solicitações de streaming, a ID de rastreamento é enviada como um evento SSE separado (data: {"trace_id": "tr-..."}) perto do final do fluxo. Em seguida, use o MLflow get_trace para recuperar o rastreamento completo usando a ID de rastreamento.

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

Agentes no Serviço de Modelo

Use o exemplo a seguir para agentes herdados hospedados no Serviço de Modelo seguindo a interface ResponsesAgent. Você pode usar um token OAuth do Databricks ou um PAT (Token de Acesso Pessoal) para consultar agentes hospedados no Serviço de Modelo. A chamada à API REST é equivalente a:

• Usando o Cliente do Databricks OpenAI com responses.create.
• Enviando uma solicitação POST para a URL do ponto de extremidade específico (por exemplo: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Para obter mais informações, consulte a página Serviço de Modelo do ponto de extremidade e a Documentação do Serviço de Modelo.

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, poderá adicioná-los ao corpo da solicitação:

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 as interfaces herdadas ChatAgent ou ChatModel. Isso é equivalente a:

• Usando o Cliente do Databricks OpenAI com chat.completions.create.
• Enviando uma solicitação POST para a URL do ponto de extremidade específico (por exemplo: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Para obter mais informações, consulte a página Serviço de Modelo do ponto de extremidade e a Documentação do Serviço de Modelo.

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, poderá adicioná-los ao corpo da solicitação:

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

Funções de IA: ai_query

Você pode usar ai_query para consultar um agente implantado hospedado no modelo que serve usando SQL. Consulte ai_query a função para sintaxe SQL e definições de parâmetro.

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

Próximas Etapas 

Monitorar aplicações de GenAI generativa na produção