はじめに

やりたいこと

AutoGen v0.2 のメッセージのやり取りの内容を自前でハンドリングしたいことってありますよね？例えば自前でどこかにログを保存しておきたいとき、例えば Streamlit などのチャット UI を作れるフレームワークにチャットのデータを逐次流したいとき。

そんな要望をかなえる、簡単な方法を紹介します。

前提条件

AutoGen v0.2 の使える環境は事前にセットアップしましょう。

そろそろ、AutoGen v0.4 もだいぶまとまりつつあるので、今後も新機能の追加が継続していく AutoGen v0.4 を選ぶ選択も検討してみるとよいかもしれません。

ちなみに、AutoGen v0.4 だと今回紹介する手法は使えず、メッセージハンドラを利用する必要があります。
https://microsoft.github.io/autogen/stable//user-guide/core-user-guide/framework/message-and-communication.html#message-handlers

サンプルコード

ここにおいてあります。

手順

まずはひな型となるサンプルコード

今回の手順では、以下のサンプルコードを使って独自の出力を作りこんでいきます。

python nofunc.py

from dotenv import load_dotenv
load_dotenv()

import os
from autogen import AssistantAgent, UserProxyAgent

# LLM への接続設定を環境変数から読み込む。キャッシュは無効化して繰り返し実行時に同じ結果にならないようにする
llm_config = {
    "cache_seed": None,
    "config_list": [
        {
            "model": os.getenv("DEPLOYMENT_NAME"),
            "api_type": "azure",
            "api_key": os.getenv("API_KEY"),
            "base_url": os.getenv("API_ENDPOINT"),
            "api_version": os.getenv("API_VERSION"),
        }
    ]
}

# LLM でタスク解決を行うためのアシスタントエージェントを定義
assistant = AssistantAgent("assistant", llm_config=llm_config)

# ユーザーの代理となるエージェント (user_proxy) を定義
user_proxy = UserProxyAgent("user_proxy",
                            code_execution_config=False,
                            human_input_mode="NEVER",
                            is_termination_msg=lambda msg: msg.get("content") is not None
                            and "TERMINATE" in msg["content"]
                            )

# user_proxy からチャットを開始する。ユーザーからの依頼となるメッセージもここで設定
user_proxy.initiate_chat(
    assistant,
    message="Microsoft の株価について、ポジティブに笑えるジョークを考えてください。一つジョークを考えおわったら TERMINATE と出力してください。",
)

デフォルトの出力はこんな感じ。

user_proxy (to assistant):

Microsoft の株価について、ポジティブに笑えるジョークを考えてください。一つジョークを考えおわったら TERMINATE と出力してください。

--------------------------------------------------------------------------------
assistant (to user_proxy):

なぜ Microsoft の株価は朝のコーヒーのようなものなのか知っていますか？

どちらも見ている間に、どんどん上がっていくんですよ！

TERMINATE

--------------------------------------------------------------------------------

追加するコード

エージェント定義用のクラスである、ConvasableAgent に、他のエージェントに返信を返す際に呼び出されるカスタム関数を追加するためのメソッドが register_reply として用意されています。

今回は、下記の関数でメッセージをコンソールに単純に出力するだけの機能を追加してみます。

# メッセージを自前でハンドリングするための関数を定義
def print_messages(recipient, messages, sender, config): 
    if "callback" in config and  config["callback"] is not None:
        callback = config["callback"]
        callback(sender, recipient, messages[-1])
    print("======================================")
    print(f"Messages sent to: {recipient.name} | num messages: {len(messages)}")
    print(messages[-1])
    print("======================================")
    print("\n")
    return False, None  # エージェントの通信フローが続行されることを保証するために必要

カスタム返信用の関数の用意ができたら、先ほどのサンプルコードのそれぞれのエージェントに対して、上記の関数を登録します。

# LLM でタスク解決を行うためのアシスタントエージェントを定義
assistant = AssistantAgent("assistant", llm_config=llm_config)
assistant.register_reply(
    [Agent, None],
    reply_func=print_messages, 
    config={"callback": None},
)

# ユーザーの代理となるエージェント (user_proxy) を定義
user_proxy = UserProxyAgent("user_proxy",
                            code_execution_config=False,
                            human_input_mode="NEVER",
                            is_termination_msg=lambda msg: msg.get("content") is not None
                            and "TERMINATE" in msg["content"]
                            )
user_proxy.register_reply(
    [Agent, None],
    reply_func=print_messages, 
    config={"callback": None},
)

完成形

その他、足りない import を足したり、AutoGen のデフォルトのコンソール出力を無効化した完成形のコードがこちら。

python withfunc.py

from dotenv import load_dotenv
load_dotenv()

import os
from autogen import AssistantAgent, UserProxyAgent,Agent

# LLM への接続設定を環境変数から読み込む。キャッシュは無効化して繰り返し実行時に同じ結果にならないようにする
llm_config = {
    "cache_seed": None,
    "config_list": [
        {
            "model": os.getenv("DEPLOYMENT_NAME"),
            "api_type": "azure",
            "api_key": os.getenv("API_KEY"),
            "base_url": os.getenv("API_ENDPOINT"),
            "api_version": os.getenv("API_VERSION"),
        }
    ]
}

# メッセージを自前でハンドリングするための関数を定義
def print_messages(recipient, messages, sender, config): 
    if "callback" in config and  config["callback"] is not None:
        callback = config["callback"]
        callback(sender, recipient, messages[-1])
    print("======================================")
    print(f"Messages sent to: {recipient.name} | num messages: {len(messages)}")
    print(messages[-1])
    print("======================================")
    print("\n")
    return False, None  # エージェントの通信フローが続行されることを保証するために必要

# LLM でタスク解決を行うためのアシスタントエージェントを定義
assistant = AssistantAgent("assistant", llm_config=llm_config)
assistant.register_reply(
    [Agent, None],
    reply_func=print_messages, 
    config={"callback": None},
)

# ユーザーの代理となるエージェント (user_proxy) を定義
user_proxy = UserProxyAgent("user_proxy",
                            code_execution_config=False,
                            human_input_mode="NEVER",
                            is_termination_msg=lambda msg: msg.get("content") is not None
                            and "TERMINATE" in msg["content"]
                            )
user_proxy.register_reply(
    [Agent, None],
    reply_func=print_messages, 
    config={"callback": None},
)

# user_proxy からチャットを開始する。ユーザーからの依頼となるメッセージもここで設定
user_proxy.initiate_chat(
    assistant,
    message="Microsoft の株価について、ポジティブに笑えるジョークを考えてください。一つジョークを考えおわったら TERMINATE と出力してください。",
    silent=True,  # silent=True にすると、autogen が出力するメッセージは表示されない
)

こんな感じでコンソール出力が自分で作った関数のものに変わっているのがわかります。

======================================
Messages sent to: assistant | num messages: 1
{'content': 'Microsoft の株価について、ポジティブに笑えるジョークを考えてください。一つジョークを考えおわったら TERMINATE と出力してください。', 'role': 'user', 'name': 'user_proxy'} 
======================================


======================================
Messages sent to: user_proxy | num messages: 2
{'content': 'Microsoftの株価が上がり過ぎた時、投資家が言った: 「この株、まるでクリップアートみたいだ、ヒントを求めなくてもすぐに役立つ！」 TERMINATE', 'role': 'user', 'name': 'assistant'}
======================================

まとめ

AutoGen v0.2 のエージェントのやり取りをカスタム関数で横取りする方法を紹介しました！今回のサンプルコードは単純な print の例でしたが、例えばこの関数の中で Streamlit の Chat UI にデータを渡すことで、おしゃれな UI を作ることも出来たりしますね！

超小ネタな感じでしたが、ぜひ活用いただけたら幸いです！

参考文献