Add functional termination condition (#6398)

Use an expression for termination condition check. This works well
especially with structured messages.

Author: ekzhu

python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py

@@ -5,6 +5,7 @@
 
 from ._terminations import (
     ExternalTermination,
+    FunctionalTermination,
     FunctionCallTermination,
     HandoffTermination,
     MaxMessageTermination,
@@ -27,4 +28,5 @@
     "SourceMatchTermination",
     "TextMessageTermination",
     "FunctionCallTermination",
+    "FunctionalTermination",
 ]

python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py

@@ -1,5 +1,6 @@
+import asyncio
 import time
-from typing import List, Sequence
+from typing import Awaitable, Callable, List, Sequence
 
 from autogen_core import Component
 from pydantic import BaseModel
@@ -154,6 +155,77 @@ def _from_config(cls, config: TextMentionTerminationConfig) -> Self:
         return cls(text=config.text)
 
 
+class FunctionalTermination(TerminationCondition):
+    """Terminate the conversation if an functional expression is met.
+
+    Args:
+        func (Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], bool] | Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[bool]]): A function that takes a sequence of messages
+            and returns True if the termination condition is met, False otherwise.
+            The function can be a callable or an async callable.
+
+    Example:
+
+        .. code-block:: python
+
+            import asyncio
+            from typing import Sequence
+
+            from autogen_agentchat.conditions import FunctionalTermination
+            from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, StopMessage
+
+
+            def expression(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> bool:
+                # Check if the last message is a stop message
+                return isinstance(messages[-1], StopMessage)
+
+
+            termination = FunctionalTermination(expression)
+
+
+            async def run() -> None:
+                messages = [
+                    StopMessage(source="agent1", content="Stop"),
+                ]
+                result = await termination(messages)
+                print(result)
+
+
+            asyncio.run(run())
+
+        .. code-block:: text
+
+            StopMessage(source="FunctionalTermination", content="Functional termination condition met")
+
+    """
+
+    def __init__(
+        self,
+        func: Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], bool]
+        | Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[bool]],
+    ) -> None:
+        self._func = func
+        self._terminated = False
+
+    @property
+    def terminated(self) -> bool:
+        return self._terminated
+
+    async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
+        if self._terminated:
+            raise TerminatedException("Termination condition has already been reached")
+        if asyncio.iscoroutinefunction(self._func):
+            result = await self._func(messages)
+        else:
+            result = self._func(messages)
+        if result is True:
+            self._terminated = True
+            return StopMessage(content="Functional termination condition met", source="FunctionalTermination")
+        return None
+
+    async def reset(self) -> None:
+        self._terminated = False
+
+
 class TokenUsageTerminationConfig(BaseModel):
     max_total_token: int | None
     max_prompt_token: int | None

python/packages/autogen-agentchat/tests/test_termination_condition.py

@@ -1,9 +1,11 @@
 import asyncio
+from typing import Sequence
 
 import pytest
 from autogen_agentchat.base import TerminatedException
 from autogen_agentchat.conditions import (
     ExternalTermination,
+    FunctionalTermination,
     FunctionCallTermination,
     HandoffTermination,
     MaxMessageTermination,
@@ -15,13 +17,17 @@
     TokenUsageTermination,
 )
 from autogen_agentchat.messages import (
+    BaseAgentEvent,
+    BaseChatMessage,
     HandoffMessage,
     StopMessage,
+    StructuredMessage,
     TextMessage,
     ToolCallExecutionEvent,
     UserInputRequestedEvent,
 )
 from autogen_core.models import FunctionExecutionResult, RequestUsage
+from pydantic import BaseModel
 
 
 @pytest.mark.asyncio
@@ -375,3 +381,53 @@ async def test_function_call_termination() -> None:
     )
     assert not termination.terminated
     await termination.reset()
+
+
+@pytest.mark.asyncio
+async def test_functional_termination() -> None:
+    async def async_termination_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> bool:
+        if len(messages) < 1:
+            return False
+        if isinstance(messages[-1], TextMessage):
+            return messages[-1].content == "stop"
+        return False
+
+    termination = FunctionalTermination(async_termination_func)
+    assert await termination([]) is None
+    await termination.reset()
+
+    assert await termination([TextMessage(content="Hello", source="user")]) is None
+    await termination.reset()
+
+    assert await termination([TextMessage(content="stop", source="user")]) is not None
+    assert termination.terminated
+    await termination.reset()
+
+    assert await termination([TextMessage(content="Hello", source="user")]) is None
+
+    class TestContentType(BaseModel):
+        content: str
+        data: str
+
+    def sync_termination_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> bool:
+        if len(messages) < 1:
+            return False
+        last_message = messages[-1]
+        if isinstance(last_message, StructuredMessage) and isinstance(last_message.content, TestContentType):  # type: ignore[reportUnknownMemberType]
+            return last_message.content.data == "stop"
+        return False
+
+    termination = FunctionalTermination(sync_termination_func)
+    assert await termination([]) is None
+    await termination.reset()
+    assert await termination([TextMessage(content="Hello", source="user")]) is None
+    await termination.reset()
+    assert (
+        await termination(
+            [StructuredMessage[TestContentType](content=TestContentType(content="1", data="stop"), source="user")]
+        )
+        is not None
+    )
+    assert termination.terminated
+    await termination.reset()
+    assert await termination([TextMessage(content="Hello", source="user")]) is None

python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb

@@ -16,6 +16,7 @@
     "- {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`: A team that runs a group chat with participants taking turns in a round-robin fashion (covered on this page). [Tutorial](#creating-a-team) \n",
     "- {py:class}`~autogen_agentchat.teams.SelectorGroupChat`: A team that selects the next speaker using a ChatCompletion model after each message. [Tutorial](../selector-group-chat.ipynb)\n",
     "- {py:class}`~autogen_agentchat.teams.MagenticOneGroupChat`: A  generalist multi-agent system for solving open-ended web and file-based tasks across a variety of domains. [Tutorial](../magentic-one.md) \n",
+    "- {py:class}`~autogen_agentchat.teams.Swarm`: A team that uses {py:class}`~autogen_agentchat.messages.HandoffMessage` to signal transitions between agents. [Tutorial](../swarm.ipynb)\n",
     "\n",
     "```{note}\n",
     "\n",

python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb

@@ -40,7 +40,8 @@
                 "7. {py:class}`~autogen_agentchat.conditions.ExternalTermination`: Enables programmatic control of termination from outside the run. This is useful for UI integration (e.g., \"Stop\" buttons in chat interfaces).\n",
                 "8. {py:class}`~autogen_agentchat.conditions.StopMessageTermination`: Stops when a {py:class}`~autogen_agentchat.messages.StopMessage` is produced by an agent.\n",
                 "9. {py:class}`~autogen_agentchat.conditions.TextMessageTermination`: Stops when a {py:class}`~autogen_agentchat.messages.TextMessage` is produced by an agent.\n",
-                "10. {py:class}`~autogen_agentchat.conditions.FunctionCallTermination`: Stops when a {py:class}`~autogen_agentchat.messages.ToolCallExecutionEvent` containing a {py:class}`~autogen_core.models.FunctionExecutionResult` with a matching name is produced by an agent."
+                "10. {py:class}`~autogen_agentchat.conditions.FunctionCallTermination`: Stops when a {py:class}`~autogen_agentchat.messages.ToolCallExecutionEvent` containing a {py:class}`~autogen_core.models.FunctionExecutionResult` with a matching name is produced by an agent.\n",
+                "11. {py:class}`~autogen_agentchat.conditions.FunctionalTermination`: Stop when a function expression is evaluated to `True` on the last delta sequence of messages. This is useful for quickly create custom termination conditions that are not covered by the built-in ones."
             ]
         },
         {
@@ -510,7 +511,7 @@
             "name": "python",
             "nbconvert_exporter": "python",
             "pygments_lexer": "ipython3",
-            "version": "3.12.9"
+            "version": "3.12.3"
         }
     },
     "nbformat": 4,