Abfragen eines bereitgestellten Mosaic-KI-Agents

Erfahren Sie, wie Sie Anforderungen an Agents senden, die auf einem Model Serving-Endpunkt bereitgestellt wurden. Databricks bietet mehrere Abfragemethoden für unterschiedliche Anwendungsfälle und Integrationsanforderungen.

Informationen zum Bereitstellen von Agents finden Sie unter Bereitstellen eines Agents für generative KI-Anwendungen.

Wählen Sie den Abfrageansatz aus, der am besten zu Ihrem Anwendungsfall passt:

Databricks empfiehlt den Databricks OpenAI Client für neue Anwendungen. Wählen Sie die REST-API aus, wenn sie in Plattformen integriert wird, die openAI-kompatible Endpunkte erwarten.

Databricks OpenAI Client (empfohlen)

Databricks empfiehlt, den Databricks OpenAI-Client zum Abfragen eines bereitgestellten Agents zu verwenden. Abhängig von der API Ihres bereitgestellten Agents verwenden Sie entweder den Antwort- oder Chatabschlussclient:

Endpunkte des Response Agents

Verwenden Sie das folgende Beispiel für Agents, die mit der ResponsesAgent-Schnittstelle erstellt wurden. Dies ist der empfohlene Ansatz für die Erstellung von Agents.

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)

Wenn Sie custom_inputs oder databricks_options übergeben möchten, können Sie sie mit dem Parameter extra_body hinzufügen.

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)

ChatAgent- oder ChatModel-Endpunkte

Verwenden Sie das folgende Beispiel für Agents, die mit älteren ChatAgent- oder ChatModel-Schnittstellen erstellt wurden, die weiterhin unterstützt, aber nicht für neue Agents empfohlen werden.

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)

Wenn Sie custom_inputs oder databricks_options übergeben möchten, können Sie sie mit dem Parameter extra_body hinzufügen.

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)

MLflow-Bereitstellungsclient

Verwenden Sie den MLflow-Bereitstellungsclient, wenn Sie in vorhandenen MLflow-Workflows und -Pipelines arbeiten. Dieser Ansatz lässt sich natürlich in die MLflow-Tracking- und Experimentverwaltung integrieren.

Die folgenden Beispiele zeigen, wie Sie einen Agent mithilfe des MLflow-Bereitstellungsclients abfragen. Für neue Anwendungen empfiehlt Databricks die Verwendung des Databricks OpenAI-Clients für seine erweiterten Features und die native Integration.

Abhängig von der API Ihres bereitgestellten Agents verwenden Sie entweder das Format "ResponsesAgent" oder "ChatAgent":

Endpunkte des Response Agents

Verwenden Sie das folgende Beispiel für Agents, die mit der ResponsesAgent-Schnittstelle erstellt wurden. Dies ist der empfohlene Ansatz für die Erstellung von Agents.

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)

ChatAgent- oder ChatModel-Endpunkte

Verwenden Sie dies für Agents, die mit älteren ChatAgent- oder ChatModel-Schnittstellen erstellt wurden, die weiterhin unterstützt, aber nicht für neue Agents empfohlen werden.

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() und client.predict_stream() rufen die von Ihnen beim Erstellen des Agents definierten Agentfunktionen auf. Siehe Streamingantworten.

REST API

Die Databricks-REST-API stellt Endpunkte für Modelle bereit, die openAI-kompatibel sind. Auf diese Weise können Sie Databricks-Agents verwenden, um Anwendungen zu bedienen, die OpenAI-Schnittstellen erfordern.

Dieser Ansatz ist ideal für:

• Sprachunabhängige Anwendungen, die HTTP-Anforderungen verwenden
• Integration mit Drittanbieterplattformen, die openAI-kompatible APIs erwarten
• Migrieren von OpenAI zu Databricks mit minimalen Codeänderungen

Authentifizieren mit der REST-API mithilfe eines Databricks OAuth-Tokens oder eines persönlichen Zugriffstokens (PERSONAL Access Token, PAT). In den folgenden Beispielen wird ein Databricks OAuth-Token verwendet. Weitere Optionen und Informationen finden Sie in der Databricks-Authentifizierungsdokumentation .

Endpunkte des Response Agents

Verwenden Sie das folgende Beispiel für Agents, die mit der ResponsesAgent-Schnittstelle erstellt wurden. Dies ist der empfohlene Ansatz für die Erstellung von Agents. DER REST-API-Aufruf entspricht:

• Verwendung des Databricks OpenAI-Clients mit responses.create.
• Senden einer POST-Anforderung an die URL des bestimmten Endpunkts (z. B. https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Weitere Details finden Sie auf der Seite "Modellbereitstellung" Ihres Endpunkts und in der Dokumentation zur Modellbereitstellung.

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

Wenn Sie custom_inputs oder databricks_options übergeben möchten, können Sie sie mit dem Parameter extra_body hinzufügen.

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

ChatAgent- oder ChatModel-Endpunkte

Verwenden Sie dies für Agents, die mit älteren ChatAgent- oder ChatModel-Schnittstellen erstellt wurden, die weiterhin unterstützt, aber nicht für neue Agents empfohlen werden. Dies entspricht:

• Verwendung des Databricks OpenAI-Clients mit chat.completions.create.
• Senden einer POST-Anforderung an die URL des bestimmten Endpunkts (z. B. https://<host.databricks.com>/serving-endpoints/\<model-name\>/invocations). Weitere Details finden Sie auf der Seite "Modellbereitstellung" Ihres Endpunkts und in der Dokumentation zur Modellbereitstellung.

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

Wenn Sie custom_inputs oder databricks_options übergeben möchten, können Sie sie mit dem Parameter extra_body hinzufügen.

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

KI-Funktionen: ai_query

Sie können einen bereitgestellten ai_query KI-Agenten mithilfe von SQL abfragen. Siehe ai_query Funktion für SQL-Syntax und Parameterdefinitionen.

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

Nächste Schritte

Überwachen von GenAI in der Produktion