Query agents deployados em Azure Databricks

A forma como você faz consultas a um agente depende do servidor de agentes que o hospeda. Encontre seu corretor na tabela a seguir, depois siga a seção correspondente. Para saber mais sobre servidores agentes, veja Agent Server.

Agentes hospedados em aplicativos Databricks requerem um token Databricks OAuth. Tokens de acesso pessoal não funcionam para aplicativos Databricks. Para gerar tokens OAuth a partir de um script, um principal de serviço, outro aplicativo ou um notebook, veja Conectar a uma API Databricks usando autenticação de token.

Consultar um agente que usa DurableAgentServer

Agentes que usam DurableAgentServer servem a API de invocação. Os agentes que você cria com a CLI Agent Bricks usam DurableAgentServer, e agentbricks deploy os implantam no Agent Runtime como um aplicativo chamado agent-bricks-<name>. Cada requisição para a API inicia uma execução do seu agente, chamada de invocação.

Corpo da solicitação

O corpo da solicitação para POST /api/invocations aceita os seguintes campos. O servidor rejeita solicitações que contenham outros campos.

O manipulador do seu agente define a estrutura de input. Agentes gerados a partir dos templates CLI leem os seguintes campos quando input é um objeto:

Para permitir que o agente se lembre do que aprendeu sobre um usuário ao longo das sessões, passe o ID do usuário como actor:

{
  "id": "550e8400-e29b-41d4-a716-446655440000",
  "session_id": "support-case-123",
  "input": {
    "messages": [{ "role": "user", "content": "What does Databricks do?" }],
    "actor": "user-42"
  }
}

Bricks do Agente

Para testar um agente implantado do seu terminal, use agentbricks endpoint invoke. O comando busca o app e autentica com seu perfil CLI.

agentbricks --profile <profile> endpoint invoke agent-bricks-<name> \
  --path /api/invocations \
  --json "{\"id\":\"$(uuidgen)\",\"session_id\":\"$(uuidgen)\",\"input\":[{\"role\":\"user\",\"content\":\"Hello\"}]}"

Para transmitir a resposta, adicione "stream":true ao corpo do JSON e passe --sse. Para testar o agente enquanto ele roda localmente com agentbricks dev, substitua o nome do app por --url http://localhost:8000.

API REST

1. Obtenha a URL do aplicativo. O campo URL na saída é a URL base para a API de invocações.

   agentbricks --profile <profile> deployments get agent-bricks-<name>
   

2. Obtenha um token OAuth para o seu perfil. A saída contém o token no access_token campo.

   databricks auth token --profile <profile>
   

3. Enviar uma solicitação:

   curl --request POST \
     --url <app-url>/api/invocations \
     --header 'Authorization: Bearer <OAuth token>' \
     --header 'content-type: application/json' \
     --data '{
       "id": "550e8400-e29b-41d4-a716-446655440000",
       "session_id": "support-case-123",
       "input": [{ "role": "user", "content": "What does Databricks do?" }]
     }'
   

A resposta contém a invocação id, seu respectivo status e um campo output com o valor que o agente retorna. O servidor não requer um esquema de saída: seu handler pode devolver qualquer valor serializável em JSON. Agentes gerados a partir dos templates CLI retornam um objeto com os seguintes campos:

• output: as mensagens que o agente produziu nesta invocação.
• status: completed, ou interrupted se o agente fez uma pausa para entrada humana.

Python

O exemplo a seguir usa o SDK Databricks para buscar a URL do app e gerar um token OAuth, e então chama a API de invocações. O WorkspaceClient deve usar autenticação OAuth.

import uuid

import requests
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
app_url = w.apps.get("agent-bricks-<name>").url
session_id = str(uuid.uuid4())

response = requests.post(
    f"{app_url}/api/invocations",
    headers=w.config.authenticate(),
    json={
        "id": str(uuid.uuid4()),
        "session_id": session_id,
        "input": [{"role": "user", "content": "What does Databricks do?"}],
    },
)
response.raise_for_status()
print(response.json()["output"])

Para continuar a conversa, envie a próxima mensagem com a mesma session_id e uma nova id.

Transmitir, executar em segundo plano e reconectar

• Stream: Definir "stream": true. A resposta é um fluxo SSE que inclui eventos run.started e run.failed (ou delta), além dos eventos que seu agente emite, como eventos run.completed com texto transmitido. Cada evento tem um ID.
• Execute em segundo plano: Defina "background": true. O servidor retorna uma 202 resposta com um status_url. Consulte repetidamente GET /api/invocations/<id> até que o status seja completed. Se você também definir "stream": true, a resposta inclui um events_url de onde você pode ler eventos.
• Reconectar: Se um stream desconectar, chame GET /api/invocations/<id>/events?after=<event-id> com o ID do último evento que você recebeu.

O exemplo a seguir em Python transmite uma resposta:

with requests.post(
    f"{app_url}/api/invocations",
    headers=w.config.authenticate(),
    json={
        "id": str(uuid.uuid4()),
        "session_id": session_id,
        "input": [{"role": "user", "content": "Summarize our last conversation."}],
        "stream": True,
    },
    stream=True,
) as response:
    response.raise_for_status()
    for line in response.iter_lines(decode_unicode=True):
        if line.startswith("data: "):
            print(line[len("data: "):])

Se você implantar o agente com mais de uma instância, envie o ID da sessão em um X-Routing-Key cabeçalho para rotear todas as solicitações de uma sessão para a mesma instância.

Consultar um agente que usa o MLflow legado AgentServer

Use essa seção para os agentes que você implanta nos Aplicativos do Databricks com o servidor de agentes herdado: o MLflow AgentServer ou LongRunningAgentServer, com a interface ResponsesAgent. Esses agentes atendem à API OpenAI Responses em /responses.

LongRunningAgentServer atende à mesma API, então os exemplos a seguir também se aplicam a ela. Também suporta execuções em segundo plano: defina background como true na requisição, e então recupere a resposta com GET /responses/<response-id>?stream=true&starting_after=<sequence-number>, que transmite os eventos após esse número de sequência.

cliente OpenAI da Databricks

A Databricks recomenda o cliente OpenAI da Databricks para esses agentes. Inclua o apps/ prefixo no nome do modelo.

from databricks.sdk import WorkspaceClient
from databricks_openai import DatabricksOpenAI

input_msgs = [{"role": "user", "content": "What does Databricks do?"}]
app_name = "<agent-app-name>"

# The WorkspaceClient must use OAuth authentication.
w = WorkspaceClient()
client = DatabricksOpenAI(workspace_client=w)

# Non-streaming request
response = client.responses.create(model=f"apps/{app_name}", input=input_msgs)
print(response)

# Streaming request
streaming_response = client.responses.create(
    model=f"apps/{app_name}", input=input_msgs, stream=True
)
for chunk in streaming_response:
    print(chunk)

Para passar custom_inputs, use o extra_body parâmetro:

response = client.responses.create(
    model=f"apps/{app_name}",
    input=input_msgs,
    extra_body={"custom_inputs": {"id": 5}},
)

Para obter a ID do rastreio de uma solicitação, inclua o cabeçalho x-mlflow-return-trace-id. 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)

API REST

Envie solicitações para o caminho /responses da URL do aplicativo. O corpo da requisição segue a API de Respostas do OpenAI, então você pode usar qualquer cliente HTTP ou ferramenta que a suporte.

curl --request POST \
  --url <app-url>/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "input": [{ "role": "user", "content": "hi" }],
    "stream": true
  }'

Para passar custom_inputs, adicione-os ao corpo da solicitação:

curl --request POST \
  --url <app-url>/responses \
  --header 'Authorization: Bearer <OAuth token>' \
  --header 'content-type: application/json' \
  --data '{
    "input": [{ "role": "user", "content": "hi" }],
    "custom_inputs": { "id": 5 }
  }'

Para obter o ID do rastreio, inclua o cabeçalho x-mlflow-return-trace-id: true. O corpo da resposta inclui o ID de rastreamento em um campo metadata.trace_id. Para requisições de streaming, o ID do rastreio chega como um evento SSE separado (data: {"trace_id": "tr-..."}) próximo ao final do fluxo.

Consultar um agente legado no Model Serving

Use esta seção para agentes legados implantados em endpoints de Model Serving. Você pode autenticar com um token Databricks OAuth ou um token de acesso pessoal. Para mover esses agentes para os Apps Databricks, veja Migrar um agente do Model Serving para os Apps Databricks.

cliente OpenAI da Databricks

Para agentes que usam a ResponsesAgent interface, chame responses.create com o nome do endpoint como modelo:

from databricks_openai import DatabricksOpenAI

input_msgs = [{"role": "user", "content": "What does Databricks do?"}]
endpoint = "<agent-endpoint-name>"

client = DatabricksOpenAI()

# Non-streaming request. Calls predict.
response = client.responses.create(model=endpoint, input=input_msgs)
print(response)

# Streaming request. Calls predict_stream.
streaming_response = client.responses.create(model=endpoint, input=input_msgs, stream=True)
for chunk in streaming_response:
    print(chunk)

Para agentes que usam as interfaces ChatAgent ou ChatModel herdadas, use o cliente de conclusão de chat:

from databricks.sdk import WorkspaceClient

messages = [{"role": "user", "content": "What does Databricks do?"}]
endpoint = "<agent-endpoint-name>"

client = WorkspaceClient().serving_endpoints.get_open_ai_client()
response = client.chat.completions.create(model=endpoint, messages=messages)
print(response)

Com qualquer um dos clientes, passe custom_inputs ou databricks_options através do extra_body parâmetro. Por exemplo, extra_body={"databricks_options": {"return_trace": True}} retorna o rastreamento com a resposta.

API REST

Para agentes que usam a ResponsesAgent interface, envie uma requisição para /serving-endpoints/responses com o nome do endpoint como modelo:

curl --request POST \
  --url https://<workspace-url>/serving-endpoints/responses \
  --header 'Authorization: Bearer <token>' \
  --header 'content-type: application/json' \
  --data '{
    "model": "<agent-endpoint-name>",
    "input": [{ "role": "user", "content": "hi" }],
    "stream": true
  }'

Para agentes que usam as interfaces ChatAgent ou ChatModel, envie uma solicitação para /serving-endpoints/chat/completions com uma lista messages em vez de input. Para passar custom_inputs ou databricks_options, adicione-os ao corpo da solicitação. Você também pode enviar requisições para a URL do /serving-endpoints/<agent-endpoint-name>/invocations endpoint. Veja Consultar modelos individuais por trás de um endpoint.

Playground de IA

Para conversar com um agente no Model Serving sem escrever código, abra o AI Playground e selecione o endpoint de serviço do agente. Para passar custom_inputs ao agente pelo Playground de IA, consulte Fornecer custom_inputs no Playground de IA e no aplicativo de revisão.

SQL com

Use ai_query para consultar um agente no Serviço de Modelo direto do SQL. Consulte a função ai_query para obter sintaxe e parâmetros.

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

Recursos adicionais

• Servidor de agente
• Tempo de Execução do Agente
• Configurar monitoramento de produção
• Consultar modelos de fundação e embedding: Consultar modelos de fundação, modelos de embedding e outros modelos diretamente, em vez de um agente.