Skip to content

Commit

Permalink
Merge branch 'main' into main
Browse files Browse the repository at this point in the history
  • Loading branch information
MervinPraison authored Nov 21, 2024
2 parents 4b48dd9 + 6e4609a commit 831a4d6
Show file tree
Hide file tree
Showing 3 changed files with 60 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from ._console import Console
from ._terminations import (
ExternalTermination,
HandoffTermination,
MaxMessageTermination,
StopMessageTermination,
Expand All @@ -15,5 +16,6 @@
"TokenUsageTermination",
"HandoffTermination",
"TimeoutTermination",
"ExternalTermination",
"Console",
]
Original file line number Diff line number Diff line change
Expand Up @@ -208,3 +208,45 @@ async def __call__(self, messages: Sequence[AgentMessage]) -> StopMessage | None
async def reset(self) -> None:
self._start_time = time.monotonic()
self._terminated = False


class ExternalTermination(TerminationCondition):
"""A termination condition that is externally controlled
by calling the :meth:`set` method.
Example:
.. code-block:: python
termination = ExternalTermination()
# Run the team in an asyncio task.
...
# Set the termination condition externally
termination.set()
"""

def __init__(self) -> None:
self._terminated = False
self._setted = False

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

def set(self) -> None:
self._setted = True

async def __call__(self, messages: Sequence[AgentMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
if self._setted:
self._terminated = True
return StopMessage(content="External termination requested", source="ExternalTermination")
return None

async def reset(self) -> None:
self._terminated = False
self._setted = False
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import pytest
from autogen_agentchat.messages import HandoffMessage, StopMessage, TextMessage
from autogen_agentchat.task import (
ExternalTermination,
HandoffTermination,
MaxMessageTermination,
StopMessageTermination,
Expand Down Expand Up @@ -226,3 +227,18 @@ async def test_timeout_termination() -> None:
assert await termination([TextMessage(content="Hello", source="user")]) is None
await asyncio.sleep(0.2)
assert await termination([TextMessage(content="World", source="user")]) is not None


@pytest.mark.asyncio
async def test_external_termination() -> None:
termination = ExternalTermination()

assert await termination([]) is None
assert not termination.terminated

termination.set()
assert await termination([]) is not None
assert termination.terminated

await termination.reset()
assert await termination([]) is None

0 comments on commit 831a4d6

Please sign in to comment.