Add sources field to TextMentionTermination (#5106)

This commit is contained in:
Leon De Andrade 2025-01-20 08:21:29 +01:00 committed by GitHub
parent 34bc82e24f
commit d9fd39a297
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 1 deletions

View File

@ -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")

View File

@ -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