メインコンテンツまでスキップ

AI エージェントの使用を開始する

Mosaic AI Agent Framework を使用して、初めての AI エージェントを構築します。このチュートリアルでは、次のことを行います。

  • Agent Frameworkを使用してエージェントを作成します。
  • エージェントにツールを追加します。
  • エージェントを Databricks モデルサービングエンドポイントにデプロイします。

エージェントとその他の AI 世代アプリの概念の概要については、「 AI 世代 AI アプリとは」を参照してください。

必要条件

ワークスペースで次の機能が有効になっている必要があります。

ノートブックの例

このノートブックには、最初の AI エージェントを作成してデプロイするために必要なすべてのコードが含まれています。ノートブック を Databricks ワークスペースにインポートして実行します。

Mosaic AIエージェントのデモ

Open notebook in new tab

エージェントの定義

AI エージェントは、次のもので構成されます。

  • 推論と意思決定が可能な大規模言語モデル (LLM)
  • LLM がテキストを生成するだけでなく、Python コードの実行やデータの取得など、さまざまなことを行うために使用できるツール

Databricks ノートブックで次のコードを実行して、単純なツール呼び出しエージェントを定義します。

  1. 必要なPythonパッケージをインストールします。

    Python
    %pip install -U -qqqq mlflow databricks-openai databricks-agents
    dbutils.library.restartPython()
    • mlflow: エージェント開発とエージェント トレースに使用されます。
    • databricks-openai: Databricks でホストされている LLM に接続するために使用されます。
    • databricks-agent: エージェントをパッケージ化してデプロイするために使用されます。
  2. エージェントを定義します。このコード スニペットでは、次のことを行います。

    • OpenAI クライアントを使用して Databricks モデルサービング エンドポイントに接続します。
    • autolog()を使用して MLflow トレースを有効にします。これにより、インストルメンテーションが追加され、クエリを送信したときにエージェントが何をするかを確認できます。
    • system.ai.python_execツールをエージェントに追加します。この組み込み Unity Catalog 関数を使用すると、エージェントは Python コードを実行できます。
    • プロンプトで LLM をクエリし、応答を処理する関数を定義します。
    Python
    import mlflow
    import json
    from databricks.sdk import WorkspaceClient
    from databricks_openai import UCFunctionToolkit, DatabricksFunctionClient

    # Get an OpenAI client configured to connect to Databricks model serving endpoints
    # Use this client to query the LLM
    openai_client = WorkspaceClient().serving_endpoints.get_open_ai_client()

    # Enable automatic tracing for easier debugging
    mlflow.openai.autolog()

    # Load Databricks built-in tools (Python code interpreter)
    client = DatabricksFunctionClient()
    builtin_tools = UCFunctionToolkit(function_names=["system.ai.python_exec"], client=client).tools
    for tool in builtin_tools:
    del tool["function"]["strict"]


    def call_tool(tool_name, parameters):
    if tool_name == "system__ai__python_exec":
    return DatabricksFunctionClient().execute_function("system.ai.python_exec", parameters=parameters)
    raise ValueError(f"Unknown tool: {tool_name}")

    def run_agent(prompt):
    """
    Send a user prompt to the LLM and return a list of LLM response messages
    The LLM is allowed to call the code interpreter tool, if needed, to respond to the user
    """
    result_msgs = []
    response = openai_client.chat.completions.create(
    model="databricks-claude-3-7-sonnet",
    messages=[{"role": "user", "content": prompt}],
    tools=builtin_tools,
    )
    msg = response.choices[0].message
    result_msgs.append(msg.to_dict())

    # If the model executed a tool, call it
    if msg.tool_calls:
    call = msg.tool_calls[0]
    tool_result = call_tool(call.function.name, json.loads(call.function.arguments))
    result_msgs.append({"role": "tool", "content": tool_result.value, "name": call.function.name, "tool_call_id": call.id})
    return result_msgs

エージェントをテストする

エージェントをテストするには、Pythonコードの実行が必要なプロンプトでエージェントをクエリします。

Python
answer = run_agent("What is the 100th fibonacci number?")
for message in answer:
print(f'{message["role"]}: {message["content"]}')

LLM の出力に加えて、詳細なトレース情報がノートブックに直接表示されます。これらのトレースは、低速または失敗したエージェント コールをデバッグするのに役立ちます。これらのトレースは、 mlflow.openai.autolog() .

エージェントをデプロイする

エージェントが作成されたので、それをパッケージ化して Databricks サービス エンドポイントにデプロイできます。デプロイされたエージェントに関するフィードバックの収集を開始するには、エージェントを他のユーザーと共有し、組み込みのチャットUIを使用してチャットします。

デプロイ用のエージェント コードの準備

デプロイ用にエージェント コードを準備するには、MLflow の ChatAgent インターフェイスを使用してラップします。ChatAgent インターフェイスは、Databricks にデプロイするエージェントをパッケージ化するための推奨される方法です。

  1. ChatAgent インターフェイスを実装するには、ユーザーのメッセージをエージェントに送信し、エージェントの応答を収集して ChatAgentResponses 形式で返す predict() 関数を定義する必要があります。

    Python
    import uuid
    from typing import Any, Optional

    from mlflow.pyfunc import ChatAgent
    from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatContext

    class QuickstartAgent(ChatAgent):
    def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
    # 1. Extract the last user prompt from the input messages
    prompt = messages[-1].content

    # 2. Call run_agent to get back a list of response messages
    raw_msgs = run_agent(prompt)

    # 3. Map each response message into a ChatAgentMessage and return
    # the response
    out = []
    for m in raw_msgs:
    out.append(ChatAgentMessage(
    id=uuid.uuid4().hex,
    **m
    ))

    return ChatAgentResponse(messages=out)
  2. 次のコードをノートブックに追加して、 ChatAgent クラスをテストします。

    Python
    AGENT = QuickstartAgent()
    for response_message in AGENT.predict({"messages": [{"role": "user", "content": "What's the 100th fibonacci number?"}]}).messages:
    print(f"role: {response_message.role}, content: {response_message.content}")
  3. すべてのエージェントコードを1つのファイルに結合して、ログに記録してデプロイできるようにします。

  • すべてのエージェント コードを 1 つのノートブック セルに統合します。
  • セルの上部に、エージェントをファイルに保存するための %%writefile quickstart_agent.py マジックコマンドを追加します。
  • セルの下部で、エージェントオブジェクトを使用して mlflow.models.set_model() に電話をかけます。これにより、予測を提供するときに使用するエージェント オブジェクトが MLflow に指示されます。この手順により、エージェント コードへのエントリ ポイントが効果的に構成されます。

ノートブックのセルは、次のようになります。

Python
%%writefile quickstart_agent.py

import json
import uuid
from databricks.sdk import WorkspaceClient
from databricks_openai import UCFunctionToolkit, DatabricksFunctionClient
from typing import Any, Optional

import mlflow
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatContext

# Add an mlflow.openai.autolog() call to capture traces in the serving endpoint

# Get an OpenAI client configured to talk to Databricks model serving endpoints
# We'll use this to query an LLM in our agent
openai_client = WorkspaceClient().serving_endpoints.get_open_ai_client()

# Load Databricks built-in tools (a stateless Python code interpreter tool)
client = DatabricksFunctionClient()
builtin_tools = UCFunctionToolkit(function_names=["system.ai.python_exec"], client=client).tools
for tool in builtin_tools:
del tool["function"]["strict"]


def call_tool(tool_name, parameters):
if tool_name == "system__ai__python_exec":
return DatabricksFunctionClient().execute_function("system.ai.python_exec", parameters=parameters)
raise ValueError(f"Unknown tool: {tool_name}")

def run_agent(prompt):
"""
Send a user prompt to the LLM, and return a list of LLM response messages
The LLM is allowed to call the code interpreter tool if needed, to respond to the user
"""
result_msgs = []
response = openai_client.chat.completions.create(
model="databricks-claude-3-7-sonnet",
messages=[{"role": "user", "content": prompt}],
tools=builtin_tools,
)
msg = response.choices[0].message
result_msgs.append(msg.to_dict())

# If the model executed a tool, call it
if msg.tool_calls:
call = msg.tool_calls[0]
tool_result = call_tool(call.function.name, json.loads(call.function.arguments))
result_msgs.append({"role": "tool", "content": tool_result.value, "name": call.function.name, "tool_call_id": call.id})
return result_msgs

class QuickstartAgent(ChatAgent):
def predict(
self,
messages: list[ChatAgentMessage],
context: Optional[ChatContext] = None,
custom_inputs: Optional[dict[str, Any]] = None,
) -> ChatAgentResponse:
prompt = messages[-1].content
raw_msgs = run_agent(prompt)
out = []
for m in raw_msgs:
out.append(ChatAgentMessage(
id=uuid.uuid4().hex,
**m
))

return ChatAgentResponse(messages=out)

AGENT = QuickstartAgent()
mlflow.models.set_model(AGENT)

エージェントのログを記録する

エージェントをログに記録し、 Unity Catalogに登録する これにより、エージェントとその依存関係がデプロイ用の 1 つのアーティファクトにパッケージ化されます。

Python
import mlflow
from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint
from pkg_resources import get_distribution

# Change the catalog name ("main") and schema name ("default") to register the agent to a different location
registered_model_name = "main.default.quickstart_agent"

# Specify Databricks resources that the agent needs to access.
# This step lets Databricks automatically configure authentication
# so the agent can access these resources when it's deployed.
resources = [
DatabricksServingEndpoint(endpoint_name="databricks-claude-3-7-sonnet"),
DatabricksFunction(function_name="system.ai.python_exec"),
]

mlflow.set_registry_uri("databricks-uc")
logged_agent_info = mlflow.pyfunc.log_model(
artifact_path="agent",
python_model="quickstart_agent.py",
extra_pip_requirements=[f"databricks-connect=={get_distribution('databricks-connect').version}"],
resources=resources,
registered_model_name=registered_model_name
)

エージェントをデプロイする

登録済みエージェントをサービス エンドポイントにデプロイします。

Python
from databricks import agents

deployment_info = agents.deploy(
model_name=registered_model_name, model_version=logged_agent_info.registered_model_version
)

エージェントエンドポイントが起動したら、 AI Playground を使用してチャットしたり、 関係者と共有 してフィードバックを得たりできます。

次のステップ

目標に基づいて次に進む場所を選択してください。

エージェントの品質を測定して改善する : 「エージェント評価のクイックスタート」を参照してください。

より高度なエージェントの構築 :非構造化データを使用してRAGを実行し、マルチターンの会話を処理し、エージェント評価を使用して品質を測定するエージェントを作成します。チュートリアル: 取得エージェントのビルド、評価、デプロイを参照してください。

他のフレームワークを使用してエージェントを構築する方法を学ぶ : LangGraph、pure Python、OpenAI などの一般的なライブラリを使用してエージェントを構築する方法を学びます。「ChatAgentを使用してエージェントを作成する」を参照してください