DatabricksでOpenAI Agents SDKを用いたマルチエージェントシステムの構築

Building a multi-agent system with OpenAI Agents SDK on Databricks | Mediumの翻訳です。

本書は著者が手動で翻訳したものであり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。

著者: David Huang and Qian Yu, Specialist Solutions Architects @ Databricks

AIエージェントの時代

NvidiaのJensen HuangやOpenAIのSam Altmanのような業界リーダーによると、2025年は「AIエージェントの年」と言われています。最初の4半期が終わった状態ですが、我々はすでにOpenAIのOperatorやSalesforceのAgentforceのような新たなAIエージェントが市場に投入されているのを目撃しています。

企業は、自分たちのデータに対する推論を行い、生産性を向上し、製品を改善するためのアクションを取ることのできるAIエージェントの開発に興味を持っています。しかし、AIエージェントの構築は簡単ではありません。これには多くの場合、LLMエンドポイントに対して繰り返し、再起的に呼び出しを行うカスタムのワークフローの構築が含まれます。この課題に対応するために、このプロセスの簡素化の助けとなるオープンソースのオーケストレーションフレームワークは数多く開発されました。

LangGraph、LlamaIndex、AutoGen、CrewAIのようなフレームワークは全てそれぞれがユニークな強さと弱さを持っており、AI開発者コミュニティで話題となっています。また、OpenAIは最近自身のAgents SDKをリリースしました - これは、マルチエージェントワークフローを構築するための軽量ですがパワフルなフレームワークです。

Databricksはフレームワークに対して不可知論のスタンスです。我々のゴールは、我々のプラットフォームを企業のユースケースのためのAIエージェントを構築、デプロイするためのベストな場所とすることです。このブログ記事では、Databricksデータインテリジェンスプラットフォームで保険ポリシーのQ&Aエージェントを構築、デプロイするために、OpenAIのAgents SDKとDatabricksの機能をどのように活用するのかをデモンストレーションします。

AIエージェントの構築

エージェントシステムの構築においては、いくつかの重要なコンポーネントが存在します:

マルチエージェントワークフローの例

• (複数の)LLM: 全てを実現するキーテクノロジーです。ユーザーのリクエストに対応したり、入力を処理するエージェントシステムは、さまざまな目的のために同じLLMを複数回呼び出したり、さまざまなタスクのために異なるLLMを複数回呼び出す必要があります。このチュートリアルでは、GPT-4oのようなOpenAIのモデルを使用します。
• オーケストレーター: ツール呼び出しのためにLLMと外部APIの間のやり取りの効率的な管理を支援する開発ツールです。これらは、迅速なプロトタイピングを可能にするために低レベルの実装の詳細を抽象化しつつも、スケーラビリティとモジュール製を維持します。ここで、OpenAIのAgents SDKが登場します。
• ツールと知識ベース: あなたのエージェントをあなたのデータで活用できるようにするための、エージェントの拡張です。Databricksにおいては、ツールや知識ベースを容易に作成、管理することができ、ツールとしてあなたのAIエージェントで利用できるようにします。
• 評価: あなたのAIエージェントのがプロダクションに適しているかどうかを知る唯一の方法は、厳密な評価ループのみです。Mosaic AI Agent Evaluationを用いることで、Databricksプラットフォーム上でエンドツーエンドのオンライン、オフラインの評価を実施することができます。これは、これ自身でブログ記事にする価値のある非常に重要なトピックです。このため、今のところはこのトピックは脇に置いておいて、のちのブログ記事でこのトピックを取り上げます。

以下のセクションでは、複数ターンの会話、データベースからの構造化データの検索、ベクトルデータベースからの非構造化データの取得を全てDatabricks上で行う保険会社向けのマルチエージェントAIチャットbotの構築方法をデモンストレーションします。さらに、どのように、このチャットbotをモデルとして登録して、REST APIエンドポイントとしてデプロイできるのかを説明します。

注意: このデモでは、以下の特定のバージョンのライブラリを使用します。また、ノートブック向けのDatabricksサーバレスコンピュートを使用します。さらに、この記事の最後では、ここでカバーするすべてのコードのGitHubリポジトリへのリンクを提供します。

openai-agents==0.0.7
unitycatalog-openai[databricks]==0.2.0
mlflow==2.21.0

Unity Catalogの関数を用いたツールと知識ベースへのアクセス

AIエージェントはあなたのドメインのデータにアクセスできてはじめて、あなたの企業ユースケースで有用なものとなり、これがツール呼び出しが重要である理由となります。基本的には、ツール呼び出し(あるいは関数呼び出し)は、関数(PythonやSQLの関数など)の呼び出しの使用可能な、構造化されたレスポンス(JSONオブジェクトなど)を生成できるようにLLMの出力を制御する方法です。LLM自身は実際には関数を呼び出しません。DatabrikcsでLLMの関数呼び出しがどのように動作するのかに関する詳細なドキュメントがあります。

ツール呼び出しの例

Databricksにおいては、特定のユーザーリクエストやタスクを完了するためにAIエージェントが使用できるツールよして、カスタムの関数を作成、共有することができます。例えば、保険企業では、顧客の主訴の詳細を格納するデータベースを保持していたり、顧客のポリシードキュメントを格納するベクトルデータベースを保有しているかもしれません。ここでは、DatabricksのUnity Catalog(UC)が重要な役割を演じます。

UCは、企業のデータとAI資産に対して集中管理されたガバナンスフレームワークを提供し、セキュリティ、発見可能性、再利用可能性、効率の観点で重大なメリットをもたらします。UCで管理されるあなたの企業の構造化データ(Deltaテーブル)と非構造化データ(ベクトルデータベース)は、ホストされUCで管理されるカスタム関数を通じて安全かつセキュアにアクセスすることができます(詳細はこちらをご覧ください)。このアプローチは、開発を円滑にするだけでなく、すべてのデータのやり取りが管理、保護されることを確実にします。

我々の例においては、AI保険エージェントのために以下のUC関数を作成します:

• search_claims_details_by_policy_no: ポリシー番号に基づいて顧客の主訴履歴を収集するSQL関数です。UCで管理されているので、UC関数の名前を定義する際には3レベルの名前空間に従います。
• policy_docs_vector_search: 入力クエリーに対する近似最近傍(ANN)アルゴリズムを用いてDatabricks Vector Searchのインデックスから適切な文書のチャンクを取得するDatabricks SQLのAI関数です。あるいは、Databricks AI BridgeのVectorSearchRetrieverToolを用いたツールを定義することもできます。詳細はこちらをご覧ください。

DatabricksのSQLエディタや、DatabricksノートブックでSpark SQLを用いて、これらを直接DatabricksのSQLコードとして実行することができます。

-- Create a SQL function that searches for claims details by policy_no
CREATE OR REPLACE FUNCTION catalog.schema.search_claims_details_by_policy_no (
    input_policy_no STRING COMMENT 'Policy number'
)
RETURNS TABLE
COMMENT 'Returns policy details about a customer given policy_no.'
RETURN 
SELECT *
FROM catalog.schema.claims_table
WHERE policy_no = input_policy_no
;

-- Create a SQL function that calls the vector_search() AI Function
CREATE OR REPLACE FUNCTION catalog.schema.policy_docs_vector_search (
    query STRING
    COMMENT 'The query string for searching insurance policy documentation.'
) 
RETURNS TABLE
COMMENT 'Executes a search on insurance policy documentation to retrieve text documents most relevant to the input query.' 
RETURN
SELECT
    chunked_text as page_content,
    map('doc_path', path, 'chunk_id', chunk_id) as metadata
FROM
    vector_search(
        index => catalog.schema.policy_docs_chunked_files_vs_index',
        query => query,
        num_results => 3
    )
;

関数を作成すると、指定したカタログとスキーマにおける関数タブで確認することができます。

あなたのエージェントからこれらのツールにアクセスできるようにするには、OpenAI、Anthropic、CrewAIのような前述の人気のフレームワークとインテグレーションしているUnity Catalog AI(UC-AI)ライブラリを用いることをお勧めします。執筆時点では、UC-AIはOpenAI Agents SDKとの完全なインテグレーションをリリースしていませんが、UC関数を実行するためにUC-AIクライアントを活用することは可能です。

OpenAI Agents SDKを用いてツールとしてUC関数を動作させるためには、function_toolデコレーターでラッピングされる必要があります。また、エージェントにコンテキストを注入するための方法、引き継ぎ、エージェントからの出力の型を矯正するための方法としてPydanticをサポートしています。ここでは、コンテキストとして使用される顧客IDとポリシー番号を引数するPydanticベースモデルでUserInfoクラスを作成します。OpenAIのコンテキスト管理のドキュメントによると、ここで設定するコンテキストはあなたのコードローカルでのみ利用でき、LLMには渡されません。

from unitycatalog.ai.core.databricks import (
    DatabricksFunctionClient,
    FunctionExecutionResult,
)
from agents import function_tool, RunContextWrapper
from pydantic import BaseModel

class UserInfo(BaseModel):
    cust_id: str | None = None
    policy_no: str | None = None

@function_tool
def search_claims_details_by_policy_no(wrapper: RunContextWrapper[UserInfo], policy_no: str) -> FunctionExecutionResult:
    print("[DEBUG]: the 'search_claims_details_by_policy_no' tool was called")
    wrapper.context.policy_no = policy_no
    client = DatabricksFunctionClient()
    return client.execute_function(
        function_name="ai.insurance_agent.search_claims_details_by_policy_no",
        parameters={"input_policy_no": wrapper.context.policy_no},
    )

@function_tool
def policy_docs_vector_search(query: str) -> FunctionExecutionResult:
    print("[DEBUG]: the 'policy_docs_vector_search' tool was called")
    client = DatabricksFunctionClient()
    return client.execute_function(
        function_name="ai.insurance_agent.policy_docs_vector_search",
        parameters={"query": query},
    )

OpenAI Agents SDKとMLflow Tracingによるオーケストレーション

マルチエージェントシステムは、ある問題を一緒に解決するために、プロンプト、コンテキスト、あるLLM(単一のエージェント)からのアウトプットを他のLLMに「引き継ぐ」必要のあるLLM間の会話のようなものです。この引き継ぎメカニズムは、OpenAI Agents SDKにおける主題となります。

はじめに、それぞれが固有のタスクと従うべきルーチンを持つ個々のエージェントの定義から始めます。それぞれのエージェントを強化するためのお好きなLLM、それれがアクセス権を持つツール、強制したいガードレールがあるかどうかを指定することができます。

ここでは、それぞれが前のセクションで定義したツールにアクセスする2つの個別のエージェントを定義します。はじめに、特にあなたのエージェントが自動的にOpenAIのモデルを使用するようにしたい場合には、環境変数にOpenAIのAPIキーを保存するようにしてください。そして、エージェントを定義するためにAgentクラスを使用します。以前定義されたのと同じUserInfoコンテキストオブジェクトもここでエージェントとして使用できることに注意してください。

import os
from agents import Agent
from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX

# Make sure your OpenAI API key is in the env
os.environ["OPENAI_API_KEY"] = dbutils.secrets.get(...)

# Create individual agents
claims_detail_retrieval_agent = Agent[UserInfo](
    name="Claims Details Retrieval Agent",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX}"
        "You are a claims details retrieval agent. "
        "If you are speaking to a customer, you were transferred from the triage agent. "
        "Use the following routine to support the customer. \n"
        "# Routine: \n"
        "1. Identify the last question asked by the customer. \n"
        "2. Use only the search tools to retrieve data about a claim. \n"
        "3. If you cannot answer the question, transfer back to the triage agent. \n"
    ),
    tools=[search_claims_details_by_policy_no],
    model="gpt-4o",
)

policy_qa_agent = Agent[UserInfo](
    name="Policy Q&A Agent",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX}"
        "You are an insurance policy Q&A agent. "
        "If you are speaking to a customer, you were transferred from the triage agent. "
        "Use the following routine to support the customer.\n"
        "# Routine: \n"
        "1. Identify the last question asked by the customer. \n"
        "2. Use only the search tools to answer a question about their policy. \n"
        "3. If you cannot answer the question, transfer back to the triage agent. \n"
    ),
    tools=[policy_docs_vector_search],
    model="gpt-4o",
)

instructionフィールドはエージェントへのプロンプトです。ここでは、OpenAIのRECOMMENDED_PROMPT_PREFIXオブジェクトを追加しました。これは、オプションの「引き継ぎ」プロンプトであり、我々はこれが非常にうまく動作することを発見しました。toolsフィールドは、指示で指定されたタスクを完了するためにエージェントが活用できるツールのリストを指定する場所です。modelフィールドでは、エージェントを強化するLLMを指定する場所となります。これは、アクセス権をお持ちであれば、すべてのOpenAIモデルでネイティブで動作します。

ご参考ですが、もしDatabricksでネイティブに利用できるAnthropicのClaude 3.7 Sonnet(お見逃しの場合には、こちらの発表をご覧ください)のような別のモデルを使いたい場合には、以下のようにこれをOpenAIChatCompletionsModelとして定義し、model引数を通じてエージェントに指定することができます:

from openai import AsyncOpenAI
from agents import OpenAIChatCompletionsModel

# Input the Model Serving variables
BASE_URL = "https://your-databricks-workspace.com/serving-endpoints"
API_KEY = dbutils.secrets.get(...)

# Instantiate the OpenAI Async client
client = AsyncOpenAI(
    base_url=BASE_URL,
    api_key=API_KEY,
)

# Create the custom llm object 
custom_llm = OpenAIChatCompletionsModel(
    model="databricks-claude-3-7-sonnet", 
    openai_client=client
)

ここのエージェントを定義したら、次にここのエージェントのルーターやマネージャーである「トリアージエージェント」を作成します。

triage_agent = Agent[UserInfo](
    name="Triage agent",
    instructions=(
        f"{RECOMMENDED_PROMPT_PREFIX}"
        "You are a helpful triaging agent. "
        "You can use your tools to delegate questions to other appropriate agents. "
        "If the customer does not have anymore questions, wish them a goodbye and a good rest of their day. "
    ),
    handoffs=[claims_detail_retrieval_agent, policy_qa_agent],
    model="gpt-4o",
)

この視覚化を支援するために、SDKネイティブのagent visualization関数があります:

マルチエージェントシステムのグラフ表現

おめでとうございます - マルチエージェントシステムを構築しました！どれだけうまく動作するのかを確認するためにチャットしてみましょう。

import mlflow
from agents import Runner

user_info = UserInfo(cust_id="7852", policy_no="102070455")

with mlflow.start_span(name="insurance_agent", span_type="AGENT") as span:
    print("[AGENT] Hello! How may I assist you?")
    while True:
        user_input = input("[USER]: ")
        if user_input.lower() == "exit":
            print("[AGENT]: Bye!")
            break
        if not user_input:
            continue
        try:
            result = await Runner.run(
                starting_agent=triage_agent, 
                input=user_input, 
                context=user_info
            )
            print("\n[AGENT]:", result.final_output)
        except Exception as e:
            print(f"\nError occurred: {str(e)}")

Runner.run()クラスメソッドは、starting_agentとしてトリアージエージェントをポイントします。context引数は、以前定義したのと同じPydanticベースモデルを受け取ります。このコードを実行することで、チャットbotとのチャットセッションを開始します。

以下は、ユーザーとチャットbotの間の簡単な会話です。みてわかりますように、チャットbotは明確化のための質問を行い、特定のツールを実行することができます。

チャットログ: 最初の質問とレスポンス

また、ユーザーの質問に基づいて異なるツールを使う他のエージェントに引き継ぎを行うことができます。

チャットログ: 二番目の質問とレスポンス

しかし、このやり取りからは内部で何が起こっていることを知るのは困難です - ここでトレーシングが役に立ちます。MLflow Tracingは、エージェントワークフローにおける処理についての詳細情報をキャプチャ可能にするオープンソースの観測可能性ツールです。入力、しゅつ録、リクエストののそれぞれの中間ステップに関係するメタデータを記録する手段を提供し、バグや予想しない挙動の原因を容易に特定できるようになります。また、OpenAIもネイティブなトレーシング機能を発表しました。しかし、MLflowがネイティブでサポートしており、Databricksノートブックとシームレスに連携しているので、エージェントを開発している際に即座にトレースを参照することができます。

私のチャットセッションのコードが示すように、すべての処理はMLflowが単一のスパンでのすべてのチャットセッションをキャプチャできるようにmlflow.start_span()でラップされています。アウトプットは以下のようになります:

ノートブック内のMLflowトレーシングUI

いろいろなことが起きていますが、引き継ぎとツール使用がはあ制している場所を観察するために、いくつかのトレースを詳しく見ていきましょう。左のパネルでは、キャプチャされたワークフローにおける中間処理を表現する階層構造を確認できます。右のパネルでは、システムプロンプト、入力文字列、利用可能なツール、エージェントのレスポンスのようなメタデータが表示されます。さらに、右のパネルでは、入出力や使用されたモデル、トークン数のような属性を確認することができます。

最初の既存の主訴をチェックするユーザーの質問からは、トリアージエージェントが回答でレスポンスするのではなく、transfer_to_引き継ぎ関数でレスポンスしていることを確認できます。これは、質問が主訴詳細取得エージェントに送信されたことを意味しています。引き継ぎ後は、主訴詳細取得エージェントが継続するために、ユーザーにポリシー番号を質問しています:

最初のエージェントワークフローの最後

このやりとりはトレースに表示されているように、単一のエージェントワークフローでキャプチャされています。会話のそれぞれのターンは、以降のエージェントワークフローでキャプチャされます。

以降のターンでは、以前の会話履歴を知っているトリアージエージェントが、再度主訴詳細取得エージェントに引き継ぎます。これは、Agent SDKの一部として自動で行われます。しかし、引き継ぎの際にどの情報を引き渡すのかを制御するために、エージェントの入力をフィルタリングすることができます。

ここでは、search_claims_details_by_policy_noが呼び出されて値を取得していることを確認できます:

エージェントによるツール実行の成功

ここで述べる重要なことは、OpenAI Agents SDKによって、OpenAIモデルを用いた際には自動でトレースデータが自動でキャプチャされ、OpenAI Developer Platformで即座に参照できるということです。OpenAIに自動でトレースデータを送信されたくないが、Databricksにトレースデータを保持したいセキュリティに注意を払っている企業は、2つの方法でこれを無効化することができます:

# Disabling it via an environment variable
import os
os.environ["OPENAI_AGENTS_DISABLE_TRACING"] = "1"

# Disabling it via the agent SDK
from agents import set_tracing_disabled
set_tracing_disabled(disabled=True)

これによって、OpenAIを使用している場合にすべてのトレーシングデータを無効化します。しかし、OpenAIモデルを使用したくて、Databricks MLflowトレーシングのみを使用したい場合には、Databricksモデルサービングを通じて提供される外部モデルエンドポイントとOpenAIChatCompletionsModelオブジェクトを使用すべきです。

次に、このシステムをモデルとしてデプロイし、REST APIエンドポイントとしてのDatabricksモデルサービングでサービングします。

Mosaic AI Agent Frameworkによる記録、登録、デプロイ

あなたのエージェントアプリケーションの品質を保証するためには、エージェントを使う専門家(SME)からのフィードバックを収集するための評価ループを構築することが重要です。これを行うための一つの手段がMosaic AI Agent Frameworkです。(評価自身のトピックが大きいので)ここでは評価のトピックはカバーしませんが、あなたのエージェントをAI PlaygroundとSMEがやりとりできるレビューアプリでテストできるように、どのようにあなたのエージェントをDatabricksにデプロイするのかを説明します。

チャットbotを「モデル」としてデプロイし、Databricksモデルサービングでサービングするには、いくつかの主要なステップが存在します:

1. あなたのエージェントコードをカスタム「モデル」としてmlflow.pyfunc.ChatAgentにラップします。このスキーマ仕様はエージェントのシナリオ向けに設計されており、OpenAIのChatCompletionスキーマと似たものですが、厳密な互換性があるものではありません。
2. マルチエージェントシステムをMLflowに記録し、UCに登録します。
3. REST APIエンドポイントとしてあなたのチャットボットをデプロイし、評価ループを開始するためのレビューアプリを起動するために、agents.deploy()関数を使用します。

それぞれのステップをウォークスルーしましょう。

エージェントを作成する際に常に使用すべきChatAgentインタフェースによって、開発者は複数のメッセージを返却し、ツール呼び出しエージェントの中間ステップを返却し、ツール呼び出しを確認し、マルチエージェントシナリオをサポートするエージェントを作成することができます。はじめに、Pythonクラスとしてエージェントを作成し、ChatAgentから継承し、必要なヘルパーメソッドを作成します:

import mlflow
from mlflow.pyfunc import ChatAgent

mlflow.openai.autolog()

class InsuranceChatAgent(ChatAgent):
    def __init__(self, starting_agent: Agent):
        self.starting_agent = starting_agent

    def _convert_to_input_text(selfself, messages: List[ChatAgentMessage]) -> str:
        """Extract the most recent user messages as input text"""
        for message in reversed(messages):
            if message.role == "user":
                return message.content
            return ""
    
    def _create_user_context(
            self, 
            context: Optional[ChatContext] = None, 
            custom_inputs: Optional[Dict[str, Any]] = None
        ) -> UserInfo:
        """Convert MLflow inputs to UserInfo object"""
        user_info = UserInfo()
        if context:
            conversation_id = getattr(context, "conversation_id", None)
            if conversation_id:
                user_info.conversation_id = conversation_id
                
            user_id = getattr(context, "user_id", None)
            if user_id:
                user_info.user_id = user_id
        return user_info

次に、MLflowトレーシングのデコレーターで、Runner.run()関数をラップする、カスタムのpredict()メソッドを記述します。これによって、チャットセッション全体が単一のトレースでラップされることになります。

さらに、predict_stream()メソッドは、エージェントからのレスポンスをストリーミングできるようにしてくれるジェネレーターを返却します。そして、UserInfoクラスにconversation_idとuser_idを追加しました。これらは、ツール呼び出しに干渉せずに、ユーザーのエージェントへのクエリーに基づいてエージェントの挙動を修正可能にするChatAgentネイティブなコンテキストの属性です。

class InsuranceChatAgent(ChatAgent):

    ...

    @mlflow.trace(name="insurance_chat_agent", span_type=SpanType.AGENT)
    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[Dict[str, Any]] = None
    ) -> ChatAgentResponse:
        input_text = self._convert_to_input_text(messages)
        user_info = self._create_user_context(context, custom_inputs)
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            result = loop.run_until_complete(
                Runner.run(
                    starting_agent=self.starting_agent,
                    input=input_text,
                    context=user_info,
                )
            )
        finally:
            loop.close()
        return ChatAgentResponse(
            messages=[
                ChatAgentMessage(
                    role="assistant",
                    content=result.final_output,
                    id=str(uuid4())
                )
            ]
        )

    @mlflow.trace(name="insurance_change_agent_stream", span_type=SpanType.AGENT)
    def predict_stream(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[Dict[str, Any]] = None
    ) -> Generator[ChatAgentResponse, None, None]:
        response = self.predict(messages, context, custom_inputs)
        for message in response.messages:
            yield ChatAgentChunk(delta=message)

最後に、チャットボットを初期化し、mlflow.models.set_model()で呼び出します。このステップで、記録されるモデルオブジェクトとしてエージェントコードを設定します。

AGENT = InsuranceChatAgent(starting_agent=triage_agent)
mlflow.models.set_model(AGENT)

次に行うのは、MLflowを通じたモデルの記録です。チャットbotはDatabricksのデータにアクセスする必要があるので、自動認証パススルーを有効化するために、MLflowのlog_model() APIのリソースパラメータを用いて依存するリソースを指定します。

import os
from mlflow.models.resources import (
    DatabricksFunction, 
    DatabricksServingEndpoint, 
    DatabricksVectorSearchIndex
)

resources = [
    DatabricksVectorSearchIndex("catalog.schema.policy_docs_chunked_files_vs_index"),
    DatabricksServingEndpoint("databricks-bge-large-en"),
    DatabricksFunction("catalog.schema.search_claims_details_by_policy_no"),
    DatabricksFunction("catalog.schema.policy_docs_vector_search")
]

mlflow.set_registry_uri("databricks-uc")
mlflow.set_experiment("/Users/[e-mail]/insurance_agent_experiment")
mlflow.openai.autolog()

with mlflow.start_run():
    logged_model_info = mlflow.pyfunc.log_model(
        artifact_path="insurance_chat_agent",
        python_model=os.path.join(os.getcwd(), "insurance_chat_agent.py"),
        input_example={
            "messages": [
                {
                    "role": "user",
                    "content": "hi, id like to check on my existing claims?",
                }
            ],
            "context": {"conversation_id": "123", "user_id": "123"},
        },
        pip_requirements=[
            "mlflow",
            "openai-agents",
            "unitycatalog-openai[databricks]==0.2.0",
            "pydantic",
        ],
        resources=resources
    )

完了すると、MLflowエクスペリメントのランでエージェントが正しく記録されていることを確認できるはずです。ランの中のアーティファクトタブでは、オブジェクトとして記録されているPythonコードも確認できるはずです。

記録されたエージェントコード

エージェントの登録はわかりやすいものとなっています。エージェントがモデルとしてUCに登録されると、従来のモデルやデータセット、関数と同じ方法で、UCによるガバナンスのメリットを享受することができます。

model_name = "insurance_chat_agent"
UC_MODEL_NAME = f"catalog.schema.{model_name}"

uc_registered_model_info = mlflow.register_model(
    model_uri=logged_model_info.model_uri, 
    name=UC_MODEL_NAME
)

最後に、REST APIエンドポイントとしてチャットボットをデプロイします。このためには、シンプルにMosaic AI Agent Frameworkツールキットに含まれるagents.deploy()を使うだけです。

from databricks import agents

agents.deploy(
    UC_MODEL_NAME,
    uc_registered_model_info.version,
    environment_vars={
        "OPENAI_API_KEY": "{{secrets/my_scope/my_api_key}}",
        "DATABRICKS_TOKEN": "{{secrets/my_scope/my_token}}",
    },
    tags={"endpoint_desc": "insurance_chat_agent"},
)

これによるいくつかのメリットがあります:

1. REST APIエンドポイントを他のアプリケーションや後段のタスクとインテグレーションできます。このようなインテグレーションの一つには、チャットボットとのやり取りやテストを開始できるDatabricks上のAI Playgroundです。
2. SMEからのフィードバックを記録する機能を提供するレビューアプリ。フィードバックは2つの異なるDeltaテーブルpayload_request_logsとpayload_assessment_logsに自動的に記録されます。これらの二つのペイロード記録テーブルはUCに存在しており、これらのテーブルの更新はDatabricksによって管理されます。

10分から15分後には、サービングページでエンドポイントが「Ready」状態になっていることを確認できるはずです。ページ右上にある「Use」の隣の矢印ボタンでは、AI Playgroundでデプロイしたエージェントを使用するか、バッチ推論で使用するか、クエリーするか、レビューアプリを用いるかを選択することができます。以下のようにテストを開始するためにPlaygroundのオプションを選択します:

AI Playgroundでのエージェントのテスト

あるいは、以下のようにレビューアプリで使用します。

レビューアプリでのエージェントのテスト

また、絵インドポイントで稼働中のチャットbotに直接クエリーすることもできます。こちらがいくつかの推奨の方法です。

まとめ

このブログ記事では、エージェント構築における重要なトピックをカバーしました: 呼び出し可能なUC AIツールの作成から、OpenAI Agents SDKを用いたエージェントの構築とオーケストレーション、エンドポイントとしてDatabricksモデルサービングへのチャットbotのデプロイまでです。しかし、これは堅牢なエージェントシステムの構築の最初のステップです。次に来るのは、さらに繰り返しや改善を行うためのSMEからのフィードバックの収集です。今後のブログ記事では評価の詳細をカバーするつもりです。

ここでカバーされたマテリアルに関してさらに質問がある場合には、Databricksアカウントチームに連絡することを躊躇わないでください。このチュートリアルで使用されたエンドツーエンドのコードはこちらのGitHubリポジトリにあります。

あなたのエージェント構築のジャーニーに幸運を！

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル