デプロイされた Mosaic AI エージェントのクエリ

MLflow デプロイ クライアントまたは Databricks OpenAI クライアントを使用して、デプロイされたエージェントに要求を送信する方法について説明します。

エージェントをデプロイする方法については、「 生成AI アプリケーション用のエージェントをデプロイする」を参照してください。

MLflow デプロイ クライアント (推奨)

Databricks では、MLflow デプロイ クライアントを使用してエンドポイントに対してクエリを実行することをお勧めします。MLflow デプロイ クライアントには、次の利点があります。

• オプションのカスタム入力を提供できます。
• MLflow トレースを要求できます。
• デプロイ クライアントの predict メソッドと predict_stream メソッドは、作成されたエージェントの動作と一致します。

次の例は、 MLflow デプロイ クライアントを使用してエージェントに対してクエリを実行する方法を示しています。messages の内容をエージェントに固有のクエリに置き換え、<agent-endpoint-name をエンドポイント名に置き換えます。エージェントがカスタム入力を受け入れる場合は、それらを入力Pythonディクショナリに含めます。「カスタム入力と出力」を参照してください。

Python

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>"

要求をフォーマットした後、非ストリーミング応答の場合は client.predict() を実行し、ストリーミング応答の場合は client.predict_stream() を実行します。predict() と predict_stream() は、エージェントの作成時に定義したエージェント関数を呼び出します。「 応答のストリーミング」を参照してください。

Python

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

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

Databricks OpenAI クライアント

または、Databricks OpenAI クライアントを使用して、デプロイされたエージェントに対してクエリを実行することもできます。Databricks OpenAI クライアントは、会話型チャットのユースケースのみをサポートします。つまり、メッセージの送受信のみが可能です。Databricks OpenAI クライアントを使用して、カスタム入力を含めたり、エンドポイントからのトレースを要求したりすることはできません。

次の例は、Databricks OpenAI クライアントを使用してクエリを送信する方法を示しています。messages の内容をエージェントに固有のクエリに置き換え、<agent-endpoint-name をエンドポイント名に置き換えます。

Python

from databricks.sdk import WorkspaceClient

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

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

要求をフォーマットしたら、 chat.completions.create()を実行します。ストリーミング応答のパラメーター stream=True を含めます。chat.completion.create() は、エージェントの作成時に定義した predict() または predict_stream() 関数を呼び出します。「 応答のストリーミング」を参照してください。

Python

## Run for non-streaming responses
response = client.chat.completions.create(model=endpoint, messages=messages)

## Include stream=True for streaming responses
streaming_response = client.chat.completions.create(
    model=endpoint, messages=msgs, stream=True
)

次のステップ

Agent Framework (MLflow 2) を使用してデプロイされたアプリの監視