Partilhar via


Consultar um agente implantado do Mosaic AI

Saiba como enviar solicitações para agentes implementados num endpoint de Modelo de Serviço. 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 adapta ao seu caso de uso:

Método Principais vantagens
Cliente OpenAI Databricks (Recomendado) Integração nativa, suporte completo a recursos, recursos de streaming
Cliente de implantações MLflow Padrões de MLflow existentes, pipelines de machine learning estabelecidos
API REST Compatível com OpenAI, agnóstico à linguagem, funciona com ferramentas existentes
Funções AI: ai_query Compatível com OpenAI, funciona com ferramentas existentes

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

O Databricks recomenda que você use o Databricks OpenAI Client para consultar um agente implantado. Dependendo da API do seu 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 criar 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, você pode adicioná-los com o extra_body param:

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 são suportadas, mas não recomendadas 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, você pode adicioná-los com o extra_body param:

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 MLflow

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

Os exemplos a seguir mostram como consultar um agente usando o cliente de implantação MLflow. Para novas aplicações, a Databricks recomenda o uso do Databricks OpenAI Client por seus recursos aprimorados e integração nativa.

Dependendo da API do seu 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 criar 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 herdadas do ChatAgent ou do ChatModel, que ainda são suportadas, mas não recomendadas 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 de agente que definiste quando criaste o agente. Consulte Respostas de streaming.

API REST

A API REST do Databricks fornece interfaces para modelos compatíveis com OpenAI. Isso permite que você use agentes Databricks para atender aplicativos que exigem interfaces OpenAI.

Esta abordagem é ideal para:

  • Aplicativos independentes de idioma 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 no código

Autentique-se com a API REST usando um token Databricks OAuth ou Personal Access Token (PAT). Os exemplos abaixo 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 criar agentes. A chamada à API REST é equivalente a:

  • Usando o Cliente Databricks OpenAI com responses.create.
  • Enviar uma solicitação POST para o 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 seu endpoint e na Documentação de serviço do 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, você pode adicioná-los com o extra_body param:

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 herdadas do ChatAgent ou do ChatModel, que ainda são suportadas, mas não recomendadas para novos agentes. Isto é equivalente a:

  • Usando o Cliente Databricks OpenAI com chat.completions.create.
  • Enviar uma solicitação POST para o 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 seu endpoint e na Documentação de serviço do 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, você pode adicioná-los com o extra_body param:

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 AI: ai_query

Você pode usar ai_query para consultar um agente de IA implantado usando SQL. Consulte ai_query 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óximos passos

Monitor GenAI em produção