Browse Source

Add sources field to TextMentionTermination (#5106)

tags/v0.4.3
Leon De Andrade GitHub 1 year ago
parent
commit
d9fd39a297
No known key found for this signature in database GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 1 deletions
  1. +6
    -1
      python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py
  2. +7
    -0
      python/packages/autogen-agentchat/tests/test_termination_condition.py

+ 6
- 1
python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py 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")


+ 7
- 0
python/packages/autogen-agentchat/tests/test_termination_condition.py 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


Loading…
Cancel
Save