AutoGen で特定の Agent が指定の回数コールされたら会話を終了させる TerminationCondition を実装する

本日時点で最新の AutoGen v0.6.4 での実装サンプルです。

AutoGen でマルチエージェントしているときに特定の Agent が特定の回数呼ばれたら会話を強制終了させたいって要件をいい感じに作れないかなーと考えてたのですが、よくよく考えると TerminationCondition を自作すればシンプルにできるじゃん思いついてメモしておきます。

Custom TerminationCondition の実装

TerminationCondition となる AgentCountTermination class はこんな感じ。

いい class 名思いつきませんでした (自分で考えず Gemini さんに考えさせてるだけやが...

from typing import Sequence
from typing_extensions import Self

from autogen_agentchat.messages import BaseChatMessage, BaseAgentEvent
from autogen_agentchat.messages import TextMessage, StopMessage
from autogen_agentchat.base import TerminationCondition
from pydantic import BaseModel


class AgentCountTerminationConfig(BaseModel):
    agent_name: str
    max_count: int = 1


class AgentCountTermination(TerminationCondition):
    """
    指定のエージェントが max_count 以上発話した時点で会話を終了させる。
    """
    def __init__(self, agent_name: str, max_count: int):
        self._target_source = agent_name
        self._max_count = max_count
        self._count = 0
        self._terminated = False

    @property
    def terminated(self) -> bool:
        return self._terminated

    async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
        for message in messages:
            if isinstance(message, TextMessage) and (self._target_source is None or self._target_source == message.source):
                self._count += 1
                if self._count >= self._max_count:
                    self._terminated = True
                    return StopMessage(content=f"Agent '{self._target_source}' has been called {self._count} time(s).", source="AgentCountTermination")
        
        return None

    async def reset(self) -> None:
        self._count = 0
        self._terminated = False

    def _to_config(self) -> AgentCountTerminationConfig:
        return AgentCountTerminationConfig(agent_name=self._target_source, max_count=self._max_count)

    @classmethod
    def _from_config(cls, config: AgentCountTerminationConfig) -> Self:
        return cls(agent_name=config.agent_name, max_count=config.max_count)

使い方

こんな感じで使えます。

 agent_1 = AssistantAgent(
    "AGENT1",
    description=...以下略
///


from xxxx import AgentCountTermination # xxxx はこの class のファイル置いたとこ
from autogen_agentchat.conditions import MaxMessageTermination

# 無限ループで事故らないように MaxMessageTermination も使用
# AGENT1 が3回発言 or team 内で20回会話したら終了の例
termination_condition = AgentCountTermination("AGENT1", 3) | MaxMessageTermination(20)