Agent name termination

python/packages/autogen-ext/src/autogen_ext/task/__init__.py

@@ -0,0 +1,3 @@
+from ._source_match_termination import SourceMatchTermination
+
+__all__ = ["SourceMatchTermination"]

python/packages/autogen-ext/src/autogen_ext/task/_source_match_termination.py

@@ -0,0 +1,37 @@
+from typing import Sequence, List
+
+from autogen_agentchat.base import TerminationCondition, TerminatedException
+from autogen_agentchat.messages import StopMessage, AgentMessage
+
+
+class SourceMatchTermination(TerminationCondition):
+    """Terminate the conversation after a specific agent responds.
+
+    Args:
+        agents (List[str]): List of agent names to terminate the conversation.
+
+    Raises:
+        TerminatedException: If the termination condition has already been reached.
+    """
+
+    def __init__(self, agents: List[str]) -> None:
+        self._agents = agents
+        self._terminated = False
+
+    @property
+    def terminated(self) -> bool:
+        return self._terminated
+
+    async def __call__(self, messages: Sequence[AgentMessage]) -> StopMessage | None:
+        if self._terminated:
+            raise TerminatedException("Termination condition has already been reached")
+        if not messages:
+            return None
+        last_message = messages[-1]
+        if last_message.source in self._agents:
+            self._terminated = True
+            return StopMessage(content=f"Agent '{last_message.source}' answered", source="SourceMatchTermination")
+        return None
+
+    async def reset(self) -> None:
+        self._terminated = False

python/packages/autogen-ext/tests/task/test_termination_condition.py

@@ -0,0 +1,29 @@
+import pytest
+
+from autogen_agentchat.base import TerminatedException
+from autogen_agentchat.messages import TextMessage, StopMessage
+from autogen_ext.task import SourceMatchTermination
+
+
+@pytest.mark.asyncio
+async def test_agent_name_termination() -> None:
+    termination = SourceMatchTermination(agents=["Assistant"])
+    assert await termination([]) is None
+
+    continue_messages = [
+        TextMessage(content="Hello", source="Assistant"),
+        TextMessage(content="Hello", source="user")
+    ]
+    assert await termination(continue_messages) is None
+
+    terminate_messages = [
+        TextMessage(content="Hello", source="user"),
+        TextMessage(content="Hello", source="Assistant")
+    ]
+    result = await termination(terminate_messages)
+    assert isinstance(result, StopMessage)
+    assert termination.terminated
+    with pytest.raises(TerminatedException):
+        await termination([])
+    await termination.reset()
+    assert not termination.terminated