Consultar um agente do Mosaic AI implantado

Saiba como enviar solicitações para agentes implantados em um endpoint de Model Serving. O Databricks fornece vários métodos de consulta para atender a diferentes casos de uso e necessidades de integração.

Para saber como implantar agentes, consulte Implantar um agente para aplicativos de IA generativos.

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)

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

Pontos de extremidade ResponsesAgent

Use o exemplo a seguir para agentes criados com a interface ResponsesAgent, que é a abordagem recomendada para a criação de agentes.

from databricks.sdk import WorkspaceClient

input_msgs = [{"role": "user", "content": "What does Databricks do?"}]
endpoint = "<agent-endpoint-name>" # TODO: update this with your endpoint name

w = WorkspaceClient()
client = w.serving_endpoints.get_open_ai_client()

## 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)

Pontos de extremidade ChatAgent ou ChatModel

Use o exemplo a seguir para agentes criados com interfaces ChatAgent ou ChatModel herdadas, que ainda têm suporte, mas não são recomendados para novos agentes.

from databricks.sdk import WorkspaceClient

messages = [{"role": "user", "content": "What does Databricks do?"}]
endpoint = "<agent-endpoint-name>" # TODO: update this with your endpoint name

w = WorkspaceClient()
client = w.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)

Cliente de implantações do MLflow

Use o cliente de implantações do MLflow ao trabalhar em fluxos de trabalho e pipelines existentes do MLflow. Essa abordagem se integra naturalmente ao gerenciamento de experimentos e acompanhamento do MLflow.

Os exemplos a seguir mostram como consultar um agente usando o cliente de implantação do MLflow. Para novos aplicativos, o Databricks recomenda usar o Cliente Do Databricks OpenAI para seus recursos aprimorados e integração nativa.

Dependendo da API do agente implantado, você usará o formato ResponsesAgent ou ChatAgent:

Pontos de extremidade ResponsesAgent

Use o exemplo a seguir para agentes criados com a interface ResponsesAgent, que é a abordagem recomendada para a criação de agentes.

from mlflow.deployments import get_deploy_client

client = get_deploy_client()
input_example = {
    "input": [{"role": "user", "content": "What does Databricks do?"}],
    ## Optional: Include any custom inputs
    ## "custom_inputs": {"id": 5},
    "databricks_options": {"return_trace": True},
}
endpoint = "<agent-endpoint-name>" # TODO: update this with your endpoint name

## Call predict for non-streaming responses
response = client.predict(endpoint=endpoint, inputs=input_example)

## Call predict_stream for streaming responses
streaming_response = client.predict_stream(endpoint=endpoint, inputs=input_example)

Pontos de extremidade ChatAgent ou ChatModel

Use isso para agentes criados com interfaces ChatAgent ou ChatModel herdadas, que ainda têm suporte, mas não são recomendados para novos agentes.

from mlflow.deployments import get_deploy_client

client = get_deploy_client()
input_example = {
    "messages": [{"role": "user", "content": "What does Databricks do?"}],
    ## Optional: Include any custom inputs
    ## "custom_inputs": {"id": 5},
    "databricks_options": {"return_trace": True},
}
endpoint = "<agent-endpoint-name>" # TODO: update this with your endpoint name

## Call predict for non-streaming responses
response = client.predict(endpoint=endpoint, inputs=input_example)

## Call predict_stream for streaming responses
streaming_response = client.predict_stream(endpoint=endpoint, inputs=input_example)

client.predict() e client.predict_stream() chamam as funções do agente que você definiu ao criar o agente. Consulte as respostas de Streaming.

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 ou um PAT (Token de Acesso Pessoal). Os exemplos a seguir usam um token OAuth do Databricks, consulte a Documentação de Autenticação do Databricks para obter mais opções e informações.

Pontos de extremidade ResponsesAgent

Use o exemplo a seguir para agentes criados com a interface ResponsesAgent, que é a abordagem recomendada para a criação de agentes. 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). Encontre mais detalhes na página de serviço de modelo do endpoint e na Documentação de 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 com o extra_body parâmetro:

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,
    "extra_body": {
      "custom_inputs": { "id": 5 },
      "databricks_options": { "return_trace": true }
    }
  }'

Pontos de extremidade ChatAgent ou ChatModel

Use isso para agentes criados com interfaces ChatAgent ou ChatModel herdadas, que ainda têm suporte, mas não são recomendados para novos agentes. 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). Encontre mais detalhes na página de serviço de modelo do endpoint e na Documentação de 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 com o extra_body parâmetro:

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,
    "extra_body": {
      "custom_inputs": { "id": 5 },
      "databricks_options": { "return_trace": true }
    }
  }'

Funções de IA: ai_query

Você pode usar ai_query para consultar um agente de IA implantado usando o 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