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

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 x-mlflow-return-trace-id cabeçalho 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 Model Serving, seguindo aResponsesAgent interface. 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 modelo que servem 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 em 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
  }'

Para enviar custom_inputs, adicione-os 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 Model Serving, seguindo aResponsesAgent interface. 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 de Modelagem de Modelos do ponto de extremidade e a Documentação de Modelagem de Modelos.

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 de Modelagem de Modelos do ponto de extremidade e a Documentação de Modelagem de Modelos.

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 o GenAI em produção