LlamaIndex WorkflowとMLflowによるツール呼び出しエージェントの構築

こちらのノートブックをウォークスルーします。LlamaIndexでもWorkflowというエージェントの制御フローを構築するフレームワークが導入されていたんですね。

LlamaIndex WorkflowとMLflowを使用したツールコールエージェントの構築

このインタラクティブなチュートリアルへようこそ。LlamaIndex WorkflowとそのMLflowとの統合を紹介するために設計されています。このチュートリアルは、Workflow、LlamaIndexの新しいアプローチを使用してLLMアプリケーションを設計し、MLflowで開発プロセスを管理するための実践的な学習体験を提供するノートブック形式で構成されています。

学べること

このチュートリアルの終わりまでに、以下のことができるようになります：

• LlamaIndex Workflowでツールコール機能を持つMVPエージェントアプリケーションを作成。
• MLflow Tracingでエージェントのアクションを観察。
• そのワークフローをMLflow Experimentにログ。
• モデルを再ロードして推論を実行。
• ログされたアーティファクトについて学ぶためにMLflow UIを探索。

インストール

MLflowのLlamaIndex Workflow APIとの統合は、MLflow >= 2.17.0およびLlamaIndex（コア）>= 0.11.16で利用可能です。パッケージをインストールした後、モジュールを正しく読み込むためにPythonカーネルを再起動する必要がある場合があります。

%pip install mlflow>=2.17.0 llama-index>=0.11.16 -qqqU
# Workflow utilはワークフローをHTMLとしてレンダリングするのに必要です
%pip install llama-index-utils-workflow -qqqU

お気に入りのLLMを選択

デフォルトでは、LlamaIndexはLLMおよび埋め込みモデルのソースとしてOpenAIを使用します。異なるLLMプロバイダーにサインアップする場合やローカルモデルを使用する場合は、Settingsオブジェクトを使用してそれらを設定してください。

オプション 1: OpenAI（デフォルト）

LlamaIndexはデフォルトでLLMおよび埋め込みモデルにOpenAI APIを使用します。この設定を続行するには、環境変数にAPIキーを設定するだけです。

import os

os.environ["OPENAI_API_KEY"] = "<YOUR_OPENAI_API_KEY>"

オプション 2: その他のホスト型LLM

他のホスト型LLMを使用したい場合は、

1. 選択したモデルプロバイダーの統合パッケージをダウンロードします。
2. 統合ドキュメントに記載されている必要な環境変数を設定します。
3. LLMインスタンスをインスタンス化し、グローバルSettingsオブジェクトに設定します。

以下のセルは、Databricksホスト型LLM（Llama3.3 70B instruct）を使用する例を示しています。

%pip install llama-index-llms-databricks

import os

os.environ["DATABRICKS_TOKEN"] = "<パーソナルアクセストークン>"
os.environ["DATABRICKS_SERVING_ENDPOINT"] = "https://<Databricksワークスペースのホスト名>/serving-endpoints/"

from llama_index.core import Settings
from llama_index.llms.databricks import Databricks

llm = Databricks(model="databricks-meta-llama-3-3-70b-instruct")
Settings.llm = llm

オプション 3: ローカル LLM

LlamaIndexはローカルホスト型LLMもサポートしています。セットアップ方法については、スターターチュートリアル（ローカルモデル）を参照してください。

MLflowエクスペリメントの作成

このチュートリアルをDatabricksノートブックで実行している場合は、このステップをスキップしてください。ノートブックを作成すると、自動的にMLflowエクスペリメントが設定されます。

Databricksなので以下はスキップします。

import mlflow

mlflow.set_experiment("MLflow LlamaIndex Workflow Tutorial")

ツールの定義

エージェントは tool オブジェクトを介してさまざまな関数やリソースにアクセスします。この例では、Python関数に基づいた最も単純な数学ツール add と multiply を定義します。実際のアプリケーションでは、ベクトル検索の取得、ウェブ検索、さらにはツールとして別のエージェントを呼び出すなど、任意のツールを作成できます。詳細については、ツールのドキュメントを参照してください。

一部のセルの冒頭にある # [USE IN MODEL] コメントは無視してください。これはこのチュートリアルの後のステップで使用されます！

# [USE IN MODEL]
from llama_index.core.tools import FunctionTool


def add(x: int, y: int) -> int:
    """2つの数字を足す有用な関数。"""
    return x + y


def multiply(x: int, y: int) -> int:
    """二つの数をかける有用な関数。"""
    return x * y


tools = [
    FunctionTool.from_defaults(add),
    FunctionTool.from_defaults(multiply),
]

ワークフローの定義

ワークフロープライマー

LlamaIndexワークフローはイベント駆動型のオーケストレーションフレームワークです。その核心には、ステップとイベントという2つの基本コンポーネントがあります。

• ステップ: ワークフロー内の実行単位。ステップは、Workflowベースクラスを実装するクラス内で@stepデコレータを付けたメソッドとして定義されます。
• イベント: ステップをトリガーするカスタムオブジェクト。StartEventとEndEventという2つの特別なイベントは、ワークフローの開始と終了時にディスパッチされるために予約されています。

各ステップは、その関数シグネチャを通じて入力および出力イベントを指定します。

@step
async def my_step(self, event: StartEvent) -> FooEvent:
    # このメソッドは、ワークフローの開始時にStartEventが発行されたときにトリガーされ、
    # その後、FooEventをディスパッチします。

各ステップのシグネチャと定義されたイベントに基づいて、LlamaIndexは自動的にワークフローの実行フローを構築します。

my_step関数が非同期関数として定義されていることに気付くかもしれません。LlamaIndexワークフローは非同期操作を第一級の機能として扱い、簡単な並列実行とスケーラブルなワークフローを可能にします。

ワークフローのもう一つの重要なコンポーネントはコンテキストオブジェクトです。このグローバルレジストリは、任意のステップからアクセス可能で、複数のイベントを通じて情報を渡す必要なく共有情報を定義することができます。

ワークフローとしてReActエージェントを定義する

以下のワークフロー定義は、定義したシンプルな数学ツールを利用するReActエージェントをモデル化しています。

# [USE IN MODEL]

# イベント定義
from llama_index.core.llms import ChatMessage, ChatResponse
from llama_index.core.tools import ToolOutput, ToolSelection
from llama_index.core.workflow import Event


class PrepEvent(Event):
    """新しいメッセージを処理し、チャット履歴を準備するイベント"""


class LLMInputEvent(Event):
    """LLMにreactプロンプト（チャット履歴）を入力するイベント"""

    input: list[ChatMessage]


class LLMOutputEvent(Event):
    """LLMの生成を表すイベント"""

    response: ChatResponse


class ToolCallEvent(Event):
    """ツール呼び出しをトリガーするイベント（ある場合）"""

    tool_calls: list[ToolSelection]


class ToolOutputEvent(Event):
    """ツール呼び出しの結果を処理するイベント（ある場合）"""

    output: ToolOutput

# [USE IN MODEL]

# ワークフロー定義
from llama_index.core import Settings
from llama_index.core.agent.react import ReActChatFormatter, ReActOutputParser
from llama_index.core.agent.react.types import ActionReasoningStep, ObservationReasoningStep
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.workflow import (
    Context,
    StartEvent,
    StopEvent,
    Workflow,
    step,
)


class ReActAgent(Workflow):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.tools = tools
        # エージェントがユーザーとの複数の対話を処理できるように、チャット履歴をメモリに保存します。
        self.memory = ChatMemoryBuffer.from_defaults(llm=Settings.llm)

    @step
    async def new_user_msg(self, ctx: Context, ev: StartEvent) -> PrepEvent:
        """新しいユーザーメッセージでワークフローを開始"""
        # StartEventは、ワークフローのrun()メソッドに渡されたキーを属性として持ちます。
        user_input = ev.input
        user_msg = ChatMessage(role="user", content=user_input)
        self.memory.put(user_msg)

        # 実行された推論ステップをコンテキストに保存します。開始時にクリアします。
        await ctx.set("steps", [])

        return PrepEvent()

    @step
    async def prepare_llm_prompt(self, ctx: Context, ev: PrepEvent) -> LLMInputEvent:
        """チャット履歴、ツール、および現在の推論（ある場合）を使用してreactプロンプトを準備"""
        steps = await ctx.get("steps", default=[])
        chat_history = self.memory.get()

        # 組み込みのプロンプトテンプレートを使用して、チャット履歴、ツール、および現在の推論からLLMを構築します。
        llm_input = ReActChatFormatter().format(self.tools, chat_history, current_reasoning=steps)
        return LLMInputEvent(input=llm_input)

    @step
    async def invoke_llm(self, ev: LLMInputEvent) -> LLMOutputEvent:
        """reactプロンプトでLLMを呼び出す"""
        response = await Settings.llm.achat(ev.input)
        return LLMOutputEvent(response=response)

    @step
    async def handle_llm_response(
        self, ctx: Context, ev: LLMOutputEvent
    ) -> ToolCallEvent | PrepEvent | StopEvent:
        """
        LLMの応答を解析して、要求されたツール呼び出しを抽出します。
        ツール呼び出しがない場合は、StopEventを発行して終了します。そうでない場合は、ツール呼び出しを処理するためにToolCallEventを発行します。
        """
        try:
            step = ReActOutputParser().parse(ev.response.message.content)
            (await ctx.get("steps", default=[])).append(step)

            if step.is_done:
                # 追加のツール呼び出しは不要です。StopEventを発行してワークフローを終了します。
                return StopEvent(result=step.response)
            elif isinstance(step, ActionReasoningStep):
                # ツール呼び出しがLLMから返されました。ツール呼び出しイベントをトリガーします。
                return ToolCallEvent(
                    tool_calls=[
                        ToolSelection(
                            tool_id="fake",
                            tool_name=step.action,
                            tool_kwargs=step.action_input,
                        )
                    ]
                )
        except Exception as e:
            error_step = ObservationReasoningStep(
                observation=f"推論の解析中にエラーが発生しました: {e}"
            )
            (await ctx.get("steps", default=[])).append(error_step)

        # ツール呼び出しや最終応答がない場合、再度繰り返します
        return PrepEvent()

    @step
    async def handle_tool_calls(self, ctx: Context, ev: ToolCallEvent) -> PrepEvent:
        """
        エラー処理を行いながら安全にツールを呼び出し、ツールの出力を現在の推論に追加します。その後、PrepEventを発行して、ReActプロンプトと解析の次のラウンドを準備します。
        """
        tool_calls = ev.tool_calls
        tools_by_name = {tool.metadata.get_name(): tool for tool in self.tools}

        # ツールを安全に呼び出します！
        for tool_call in tool_calls:
            if tool := tools_by_name.get(tool_call.tool_name):
                try:
                    tool_output = tool(**tool_call.tool_kwargs)
                    step = ObservationReasoningStep(observation=tool_output.content)
                except Exception as e:
                    step = ObservationReasoningStep(
                        observation=f"ツール {tool.metadata.get_name()} の呼び出し中にエラーが発生しました: {e}"
                    )
            else:
                step = ObservationReasoningStep(
                    observation=f"ツール {tool_call.tool_name} は存在しません"
                )
            (await ctx.get("steps", default=[])).append(step)

        # 次のイテレーションを準備します
        return PrepEvent()

ワークフローを視覚的に確認する

エージェントオブジェクトをインスタンス化する前に、ワークフローが期待通りに構築されているかどうかを一時停止して確認しましょう。

これを確認するために、draw_all_possible_flowsユーティリティ関数を使用してワークフローのグラフィカルな表現をレンダリングできます。

（注：レンダリングされたHTMLが空白の場合、Jupyterの安全機能が原因である可能性があります。その場合は、!jupyter trust llama_index_workflow_tutorial.ipynbでノートブックを信頼できます。詳細はJupyterのドキュメントを参照してください。）

from IPython.display import HTML
from llama_index.utils.workflow import draw_all_possible_flows

draw_all_possible_flows(ReActAgent, filename="workflow.html")

with open("workflow.html") as file:
    html_content = file.read()
HTML(html_content)

# [USE IN MODEL]
agent = ReActAgent(timeout=180)

ワークフローを実行する（トレース付き）

ワークフローの準備が整いました！しかし、その前にMLflowトレースをオンにすることを忘れないでください。これにより、エージェントの実行中の各ステップを観察し、後でレビューするために記録できます。

MLflowはLlamaIndexワークフローの自動トレースをサポートしています。これを有効にするには、mlflow.llama_index.autolog()関数を呼び出すだけです。

import mlflow

mlflow.llama_index.autolog()

# ワークフローの実行
await agent.run(input="(123 + 456) * 789はいくつ？")

'(123 + 456) * 789は456831です。'

トレースを確認する

Databricksでは自動でトレースが表示されるので、以下はスキップします。

生成されたトレースは自動的にMLflow実験に記録されます。

1. ターミナルを開き、現在のディレクトリ内で mlflow ui --port 5000 を実行します（実行したままにしておきます）。
2. ブラウザで http://127.0.0.1:5000 に移動します。
3. 実験「MLflow LlamaIndex Workflow Tutorial」を開きます。
4. 実験名のヘッダー下にある「Trace」タブに移動します。

トレースは、ワークフロー実行内の各ステップの入力、出力、およびレイテンシーなどの追加メタデータを記録します。トレースUIで次の情報を見つける簡単な演習を行いましょう。

1. 最初のLLM呼び出しで使用されたトークン数

LLM呼び出しのトークン数は、LLM呼び出しスパンの属性セクション内のusageフィールドで見つけることができます。

2. "add"ツール呼び出しの入力数

入力数x=123およびy=456は、FunctionTool.callという名前のスパンのInputsフィールドで見つけることができます。そのスパンはReActAgent.handle_tool_callsステップスパンの下にあります。

ワークフローをMLflowエクスペリメントに記録する

LlamaIndexワークフローを使用して最初のReActエージェントを構築したので、パフォーマンスを向上させるために反復的に改善および最適化することが重要です。MLflowエクスペリメントは、これらの改善を記録および管理するための理想的な場所です。

モデルスクリプトの準備

MLflowは、コードからのモデルメソッドを使用してLlamaIndexワークフローのログをサポートしており、スタンドアロンのPythonスクリプトから直接モデルを定義してログに記録することができます。このアプローチは、pickleのようなリスクが高く壊れやすいシリアル化方法を回避し、コードをモデル定義の唯一の信頼できる情報源として使用します。MLflowの環境フリーズ機能と組み合わせることで、モデルを確実に永続化する方法を提供します。

詳細については、MLflowのドキュメントを参照してください。

このノートブックからコードをコピーして、別のPythonファイルを手動で作成することもできます。しかし、便利なように、このノートブックの内容からモデルスクリプトを自動的に生成するユーティリティ関数を定義します。以下のセルを実行すると、このスクリプトが現在のディレクトリに作成され、MLflowのログに使用できるようになります。

def generate_model_script(output_path, notebook_path="llama_index_workflow_tutorial.ipynb"):
    """
    必要なライブラリのインポートとモデル定義を含む、ログに準備された.pyスクリプトを生成するユーティリティ関数。

    Args:
       output_path: .pyファイルを書き込むパス。
       notebook_path: チュートリアルノートブックのパス。
    """
    import nbformat

    with open(notebook_path, encoding="utf-8") as f:
        notebook = nbformat.read(f, as_version=4)

    # コードセルで指定されたマーカーを含むセルをフィルタリング
    merged_code = (
        "\n\n".join(
            [
                cell.source
                for cell in notebook.cells
                if cell.cell_type == "code" and cell.source.startswith("# [USE IN MODEL]")
            ]
        )
        + "\n\nimport mlflow\n\nmlflow.models.set_model(agent)"
    )

    # 出力 .py ファイルに書き込む
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(merged_code)

    print(f"モデルコードが {output_path} に保存されました")


# ノートブック名を変更した場合は `notebook_path` 引数を渡す
generate_model_script(output_path="react_agent.py")

モデルのログ

import mlflow

with mlflow.start_run(run_name="react-agent-workflow"):
    model_info = mlflow.llama_index.log_model(
        "react_agent.py",
        artifact_path="model",
        # 入力例を使用してログを記録すると、MLflowが依存関係とシグネチャ情報を正確に記録するのに役立ちます。
        input_example={"input": "(123 + 456) * 789はいくつ？"},
    )

MLflow UIの探索

ノートブックエクスペリメントを開きます。

もう一度MLflow UIを開いて、実験で追跡されている情報を確認しましょう。

1. トレースを確認したときと同様に、MLflow UIにアクセスします。
2. 実験「MLflow LlamaIndex Workflow Tutorial」を開きます。
3. 実験のRunsタブには「react-agent-workflow」という名前のランが含まれているはずです。これを開きます。
4. ランのページで、"Artifacts"タブに移動します。

Artifactsタブには、MLflowがランで保存したさまざまなファイルが表示されます。以下の画像を参照し、注釈付きのファイルを開いて、各ファイルに保存されている情報を確認してください。

推論のためのモデルの再読み込み

必要なメタデータがすべてMLflowにログされているので、異なるノートブックでモデルを読み込んだり、推論のためにデプロイしたりする際に環境の不整合を心配する必要はありません。実験結果の再現性を示すために、簡単な演習を行いましょう。

異なる環境をシミュレートするために、グローバルSettingsオブジェクトからllm設定を削除します。

from llama_index.core.llms import MockLLM

Settings.llm = MockLLM(max_tokens=1)

await agent.run(input="(123 + 456) * 789はいくつ？")

'text'

ダミーLLMが設定されているため、ワークフローは正しい出力を生成できず、単にtextを返します。

次に、mlflow.llama_index.load_model() APIを呼び出してMLflow実験からモデルを再読み込みし、ワークフローを再実行してみてください。

loaded_model = mlflow.llama_index.load_model("runs:/e83f94c199194a488e9c3d8819dcfa33/model")
await loaded_model.run(input="(123 + 456) * 789はいくつ？")

'(123 + 456) * 789は456831です。'

今回は、MLflowがログ時の元のLLM設定を自動的に復元するため、出力が正しく計算されます。

詳細を学ぶ

おめでとうございます！ 🎉 LlamaIndex WorkflowとMLflowを使用してツールコールエージェントを構築する方法を無事に学びました。

これらの高度なリソースで旅を続けましょう：

• ワークフローの品質向上: MLflow LLM評価を使用してパフォーマンスを向上させるためにワークフローを評価します。
• モデルのデプロイ: MLflowデプロイを使用して、MLflowモデルをサービングエンドポイントにデプロイします。
• さらに多くの例を探索: 公式ドキュメントでLlamaIndex Workflowの追加の例を発見します。

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル