Compartir vía

Consulta de un agente de Mosaic AI implementado

Aprenda cómo enviar solicitudes a los agentes desplegados en un endpoint de Model Serving. Databricks proporciona varios métodos de consulta para adaptarse a diferentes casos de uso y necesidades de integración.

Para obtener información sobre cómo implementar agentes, consulte Implementación de un agente para aplicaciones de IA generativas.

Seleccione el enfoque de consulta que mejor se adapte a su caso de uso:

Método Ventajas principales
Cliente openAI de Databricks (recomendado) Integración nativa, compatibilidad completa con características, funcionalidades de streaming
Cliente de implementaciones de MLflow Patrones existentes de MLflow, canalizaciones de ML ya establecidas
REST API Compatible con OpenAI, independiente del lenguaje, funciona con herramientas existentes
Funciones de IA: ai_query Compatible con OpenAI, funciona con herramientas existentes

Databricks recomienda el cliente de OpenAI de Databricks para las nuevas aplicaciones. Elija la API REST al integrar con plataformas que esperan puntos de conexión compatibles con OpenAI.

Databricks recomienda usar el cliente openAI de Databricks para consultar un agente implementado. En función de la API del agente implementado, usará las respuestas o el cliente de finalizaciones de chat:

Puntos de conexión de ResponsesAgent

Use el ejemplo siguiente para los agentes creados con la interfaz ResponsesAgent, que es el enfoque recomendado para compilar 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)

Si desea pasar en custom_inputs o databricks_options, puede agregarlos con el parámetro extra_body.

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)

Puntos de conexión de ChatAgent o ChatModel

Use el ejemplo siguiente para los agentes creados con interfaces ChatAgent o ChatModel heredadas, que siguen siendo compatibles pero no recomendadas para los nuevos 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)

Si desea pasar en custom_inputs o databricks_options, puede agregarlos con el parámetro extra_body.

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 implementaciones de MLflow

Use el cliente de implementaciones de MLflow al trabajar en flujos de trabajo y canalizaciones de MLflow existentes. Este enfoque se integra naturalmente con el seguimiento de MLflow y la administración de experimentos.

En los ejemplos siguientes se muestra cómo consultar un agente mediante el cliente de implementación de MLflow. En el caso de las nuevas aplicaciones, Databricks recomienda usar el cliente openAI de Databricks para sus características mejoradas e integración nativa.

En función de la API del agente implementado, usará el formato ResponsesAgent o ChatAgent:

Puntos de conexión de ResponsesAgent

Use el ejemplo siguiente para los agentes creados con la interfaz ResponsesAgent, que es el enfoque recomendado para compilar 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)

Puntos de conexión de ChatAgent o ChatModel

Úselo para agentes creados con interfaces de ChatAgent o ChatModel heredadas, que todavía se admiten pero no se recomiendan para los nuevos 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() y client.predict_stream() llaman a las funciones del agente que definió al crear el agente. Consulte Respuestas de streaming.

REST API

La API REST de Databricks proporciona puntos de conexión para los modelos compatibles con OpenAI. Esto le permite usar agentes de Databricks para atender aplicaciones que requieren interfaces openAI.

Este enfoque es ideal para:

  • Aplicaciones independientes del lenguaje que usan solicitudes HTTP
  • Integración con plataformas de terceros que esperan API compatibles con OpenAI
  • Migración de OpenAI a Databricks con cambios mínimos de código

Autentíquese con la API REST mediante un token de OAuth de Databricks o un token de acceso personal (PAT). En los ejemplos siguientes se usa un token de OAuth de Databricks, consulte la documentación de autenticación de Databricks para obtener más opciones e información.

Puntos de conexión de ResponsesAgent

Use el ejemplo siguiente para los agentes creados con la interfaz ResponsesAgent, que es el enfoque recomendado para compilar agentes. La llamada a la API rest es equivalente a:

  • Uso del Cliente OpenAI de Databricks con responses.create.
  • Enviar una solicitud POST a la dirección URL del punto de conexión específico (por ejemplo: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Busque más detalles en la página de implementación del modelo del punto de conexión y en la documentación de implementación 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
  }'

Si desea pasar en custom_inputs o databricks_options, puede agregarlos con el parámetro extra_body.

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

Puntos de conexión de ChatAgent o ChatModel

Úselo para agentes creados con interfaces de ChatAgent o ChatModel heredadas, que todavía se admiten pero no se recomiendan para los nuevos agentes. Esto equivale a:

  • Uso del Cliente OpenAI de Databricks con chat.completions.create.
  • Enviar una solicitud POST a la dirección URL del punto de conexión específico (por ejemplo: https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Busque más detalles en la página de implementación del modelo del punto de conexión y en la documentación de implementación 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
  }'

Si desea pasar en custom_inputs o databricks_options, puede agregarlos con el parámetro extra_body.

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

Funciones de IA: ai_query

Puede usar ai_query para consultar un agente de IA implementado mediante SQL. Consulte la función ai_query para la sintaxis SQL y las definiciones de parámetros.

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

Pasos siguientes

Supervisión de GenAI en producción

  • Last updated on