電通総研 テックブログ

電通総研が運営する技術ブログ

LLMエージェントの実行結果をPrompt flowとAzure AI Studioのトレース機能で管理する

こんにちは!Xイノベーション本部AIトランスフォーメーションセンター(AITC)の後藤です。

本記事では最近盛り上がりつつあるLLMエージェントの開発に便利なAzure AI StudioとPrompt flowを使ったトレース機能について紹介します。

はじめに

LLMエージェントは一度の実行で内部的に複数のLLMがさまざまなプロンプトで呼び出されるため、処理結果の追随が単純なLLMのチャットと比較して困難です。

このような課題に対してAzure AI Studio(Prompt flow)のトレース機能を使えば簡単にLLMの実行ログが集められると知り、今回記事にしました。

ちなみにこのようなLLMの実行結果管理ツールとしてはLangSmithというツールもあります。こちらは個人利用であればすぐに使えて便利です。今回はなるべくAzureに情報をすべて集めた方が管理しやすいのでAzure AI Studioのトレース機能を調査しました。

そもそもエージェントとは何かや、エージェントに関するより細かな内容を知りたい場合はAITCのコラムもご参照ください。

Azure AI Studio とは

トレース機能の前にAzure AI Studioの説明をします。Azure AI Studio は2024年5月にGAされたばかりのサービスです。

Azure AI Studio は、AI 開発を簡単に行える統合プラットフォームです。開発者やデータサイエンティスト向けに設計されており、AIモデルの構築、評価、デプロイをシームレスに行えます。

特徴としてはAzureでこれまで提供されてきたAI系のサービスを接続し、Azure AI Studioからそれらの機能を使用可能な点です。

また今年のMiscrosoft Buildでも発表があったようにAzure AI Studioからデプロイできるモデルも今後どんどん増えていく予定です。

これまでAzure上のLLMを用いた開発ではAzure OpenAI ServiceでOpenAIモデルのデプロイやPlayGroundでの動作確認を行い、Azure AI Searchで検索インデックスを作り、Azure Machien LearningのPrompt flowで精度検証を行い、、、のようにAzureの複数リソースをまたがって管理する必要がありました。

Azure AI Studioを用いると上記の関連サービスをAzure AI Studioを通して使えるようになります。

Azure AI Studioのトレース機能

Azure AI Studioのトレース機能とは何かは公式ドキュメントに記載があります。

トレースは、エージェント、AutoGen、取得拡張生成 (RAG) ユース ケースなどの生成 AI アプリケーションの実行プロセスを開発者が深く理解できるようにする強力なツールです。

参考:https://learn.microsoft.com/ja-jp/azure/ai-studio/how-to/develop/trace-local-sdk?tabs=python

書いてある通り、LLM関連の開発時に便利な実行プロセスの追随・可視化機能を提供してくれます。

なおAzure AI StudioはGAになりましたが、トレース機能はまだプレビュー段階の機能になります。

Azure AI Studioのトレース機能の使い方

実際にAzure AI Studioのトレース機能を試してみます。

今回はLLMとしてAzure OpenAI ServiceのAPIを使用します。

Azure AI Studioの準備

Azure AI Studio側の準備を行います。この手順は公式ドキュメントに詳細が記載されているため簡単に説明します。

参考:https://learn.microsoft.com/ja-jp/azure/ai-studio/how-to/create-azure-ai-resource

Azure AI Studioはざっくりいうとハブで複数のAzure上のAI関連リソースとの接続を定義し、ハブの中のプロジェクトという単位でハブで接続しているリソースを用いて様々な実験・開発を行います。

したがって、まずハブを作成し次に作成したハブの中にプロジェクトを作成します。


作成したプロジェクトを選択すると以下の画面になります。ここまででAzure AI studio側の準備は完了です。

ローカル環境での準備

作業用のフォルダに移動し、以下のコマンドを実行します。<>の部分は各々の環境に応じて書き換えてください。

az login
pf config set trace.destination=azureml://subscriptions/<your_subscription_id>/resourcegroups/<your_resourcegroup_name>/providers/Microsoft.MachineLearningServices/workspaces/<your_studio_project_name>

上記コマンドを実行すると自分の場合は以下のerrorが発生しました。

pf.config.set failed with MissingAzurePackage: "promptflow[azure]" is required for this functionality, please install it by running "pip install promptflow-azure" with your version.

言われた通りパッケージをインストールして再実行します。

[2024-05-30 15:56:28 +0900][promptflow.azure._restclient.flow_service_caller][INFO] - start polling until Cosmos DB setup finished...

このコマンド実行が終われば準備完了です。

トレース機能の使い方

Prompt flow(promptflow-tracing)をインストールし、以下のコードを実行するだけでLLMの実行箇所については自動で記録されます。(※ Azure OpenAI Serviceのgptモデルでしか試していません)

from promptflow.tracing import start_trace
# start a trace session, and print a url for user to check trace
start_trace(collection="trace-openai-test")

ノートブックで上記内容を実行すると以下のように出力されます。

Starting prompt flow service... Start prompt flow service on port 23333, version: 1.11.0. You can stop the prompt flow service with the following command:'pf service stop'. Alternatively, if no requests are made within 1 hours, it will automatically stop.

出力にあるようにこのコードを実行すると以下のようにローカルでブラウザからcollectionごとにトレースの確認が可能になります。


ここでいうcollectionはコード内のtrace-openai-testのことであり、実行結果はこのcollection単位でまとめられます。したがって、collectionは意味のあるまとまりの単位で都度生成した方が後から確認しやすいです。

実際にLLMを実行してどのように記録されるか確認します。上記のコードを実行後にAzure OpenAIのAPIを叩いてみます。

import os
from openai import AzureOpenAI

client = AzureOpenAI(
    api_key= os.getenv("AZURE_OPENAI_API_KEY"),
    api_version="2024-05-01-preview",
    azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
    )

response = client.chat.completions.create(
    model="gpt-4o-202405-13",
    temperature=0.7,
    max_tokens=500,
    messages=[
        {"role": "system", "content": "あなたは優秀なアシスタントです"},
        {"role": "user", "content": "夏の俳句を詠んでください"},
    ]
)

以下がcollectionの中身です。LLMの実行ごとに1行 追加されます(何度か実行した後なので複数結果があります)

それぞれの中身は以下のようになっており、ロールごとの生成結果と右上に実行時間とトークン数も記載されます。

ここまで見せたのはローカルから確認した結果です。次にAzure AI Studioにいき、作成したプロジェクトのtraceをみるとローカルと同様の内容で記録されていることがわかります。

エージェントワークフローでトレース機能を使ってみる

エージェントワークフローについては以下の記事で解説しておりますのでご参照ください。

今回は以下のような簡単なフローでトレース機能の効果を見てみます。

処理の流れとしては最初にタスクの計画を立てて、計画内容の数だけタスクを実行して最後にタスクの実行結果を踏まえて最終回答を生成する流れです。

エージェントワークフローの実装部分はLangGraphで行っています。LangGraphでフローを組んだ場合、ノードの各処理とノード間を条件付きのエッジで接続する際には関数を定義します。

トレース機能にはデコレータも用意されており、LLMの入出力だけでなく@trace デコレータが付けられた関数の入出力についても記録してくれます。

この機能はLangGraphと非常に相性が良く結果を見た際にフローの流れがわかりやすいです。

実装コード

まずはフローを定義するLangGraphの実装部分を紹介します。長くなるのでプロンプトやモデル定義部分は割愛します。

import operator
from typing import TypedDict, List, Tuple,Annotated
from langgraph.graph import StateGraph, END
from promptflow.tracing import start_trace, trace

# start a trace session, and print a url for user to check trace
start_trace(collection="trace-langgraph-workflow")

# Agent(Node)の状態を記録するためのクラス
# 基本的に各ノードにこのクラスが引数に渡される
class AgentState(TypedDict):
    input: str
    answer: str
    plan: List[str]
    past_steps: Annotated[List[Tuple], operator.add]
    exec_subtask_idx: int # 現在実行中のサブタスクのインデックス

# Graph全体を定義
workflow = StateGraph(AgentState)

# 使用するNodeを追加。名前はこの後も使うので一意である必要がある
workflow.add_node("planner", create_plan)
workflow.add_node("agent_executor", execute_action)
workflow.add_node("answer_creator", create_answer)

# エントリーポイントを定義。これが最初に呼ばれるNode
workflow.set_entry_point("planner")

# Nodeをつなぐエッジを追加
workflow.add_edge("planner", "agent_executor")
workflow.add_edge("answer_creator", END)

# 条件付きエッジを追加。reflection処理に入るか否かを判定する
workflow.add_conditional_edges(
    "agent_executor",  # つなぎ元のNode
    should_continue_execute_action,  # 条件判定関数
    {
        # 結果が"continue"ならactionにつなぐ
        "continue": "agent_executor",
        # 結果が"end"なら終了
        "end": "answer_creator",
    },
)

# 最後にworkflowをコンパイルする。これでLangChainのrunnnableな形式になる
app = workflow.compile()

上記コードでは.add_nodeでノード名とノードに対応する関数が渡されています。この関数がそのノードで実行される処理に該当します。

それらを記述したのが以下になります。@traceデコレータはこれらの関数に付けます。

@trace
def should_continue_execute_action(state):
    # 実行すべきアクションが残っているかを判定
    past_steps = state["past_steps"]

    if len(past_steps) == len(state["plan"]):
        print("end agent execution")
        return "end"
    else:
        print("continue agent execution")
        return "continue"

@trace
def create_plan(state):    
    # 回答のため実行計画を作成
    plan = planner_chain.invoke({"input": state["input"]})
    return {
        "plan": plan.steps,
        "past_steps": [],
        "exec_subtask_idx": 0,
        "reflection_try_count": 0,
        "reflection_advise": "",
        "reflection_result": False,       
    }

@trace
def execute_action(state):
    # stateからメッセージを取得
    exec_subtask_idx = state["exec_subtask_idx"]
    subtask = state["plan"][exec_subtask_idx]
    
    # サブタスクを実行
    result = agent_executor.invoke({"input": state["input"], "plan": state["plan"], "sabtask": subtask})

    return {
        "past_steps": [(subtask, result["output"])],
        "exec_subtask_idx": exec_subtask_idx + 1,
    }


@trace
def create_answer(state):
    # 最終回答を作成
    answer = answer_chain.invoke(
        {"input": state["input"], "past_steps": state["past_steps"]}
    )

    return {"answer": answer}

トレース結果を確認

上記のフローを実行してみました。ローカルで結果を確認すると以下のような表示になります。

これまでとは異なり一番左のkindがFunctionになっています。また先ほどはName部分がopenai_chatに統一されていましたが、今回は関数名です。inputとoutputに関しても関数の入出力に対応しています。

最初に実行されるcreate_plan関数の実行結果をみてみます。


Functionが外側にあり、中にLLMの実行結果が表示されています。


toolの定義についてもちゃんと確認できます。

ここまでは関数内部でLLMが使われたパターンですが、LLMが使われないパターンに関しても記録されています。

以下の画像はshould_continue_execute_action関数の実行結果です。こちらの関数はstateに蓄積された値を参照し、未実行のタスクがないかを確認する関数になります。

入力に与えた状態とその結果として出力される値が表示されており、こちらも処理フローのデバッグに大変便利です。


Azure AI Studio側の結果も見てみます。

ローカルと同様の内容があります。中身の実行結果部分に関しても正しく記録されていることがわかります。
また関数でLLMが2回呼ばれた場合もちゃんとトレースできています。

トレース機能の便利な点まとめ

Azure AI Studio(Prompt flow)のトレース機能を使ってみてわかった便利な点まとめです。

  • わずか2行という少ないコード記述量で自動でローカルとAzure AI Studio両方に記録をとってくれる
  • エージェントワークフローのような複雑な処理の実験管理が楽になる。特にLangGraphとの相性も良い
  • Azure AI Studioと連携することで結果を他人にURLで共有可能で説明がしやすい。特にtoolを使用した呼び出し部分は可視化されると格段に説明がしやすい
  • エージェントワークフローのような長い処理のデバッグに便利。途中の入出力をコピペして処理の途中から実行してみたりなども記録が残ってるとやりやすい

執筆:@goto.yuki、レビュー:@handa.kenta
Shodoで執筆されました