diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py b/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py index e33f2fcb7..7c46545e8 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py @@ -100,14 +100,16 @@ class TextMentionTermination(TerminationCondition, Component[TextMentionTerminat Args: text: The text to look for in the messages. + sources: Check only messages of the specified agents for the text to look for. """ component_config_schema = TextMentionTerminationConfig component_provider_override = "autogen_agentchat.conditions.TextMentionTermination" - def __init__(self, text: str) -> None: + def __init__(self, text: str, sources: Sequence[str] | None = None) -> None: self._text = text self._terminated = False + self._sources = sources @property def terminated(self) -> bool: @@ -117,6 +119,9 @@ class TextMentionTermination(TerminationCondition, Component[TextMentionTerminat if self._terminated: raise TerminatedException("Termination condition has already been reached") for message in messages: + if self._sources is not None and message.source not in self._sources: + continue + if isinstance(message.content, str) and self._text in message.content: self._terminated = True return StopMessage(content=f"Text '{self._text}' mentioned", source="TextMentionTermination") diff --git a/python/packages/autogen-agentchat/tests/test_termination_condition.py b/python/packages/autogen-agentchat/tests/test_termination_condition.py index cc3cd8107..68c782e0c 100644 --- a/python/packages/autogen-agentchat/tests/test_termination_condition.py +++ b/python/packages/autogen-agentchat/tests/test_termination_condition.py @@ -88,6 +88,13 @@ async def test_mention_termination() -> None: await termination([TextMessage(content="Hello", source="user"), TextMessage(content="stop", source="user")]) is not None ) + termination = TextMentionTermination("stop", sources=["agent"]) + assert await termination([TextMessage(content="stop", source="user")]) is None + await termination.reset() + assert ( + await termination([TextMessage(content="stop", source="user"), TextMessage(content="stop", source="agent")]) + is not None + ) @pytest.mark.asyncio