From 7bdd7f6162f4207338909c8fd181bbe055c6e25d Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Fri, 25 Apr 2025 15:57:36 -0700 Subject: [PATCH] Add functional termination condition (#6398) Use an expression for termination condition check. This works well especially with structured messages. --- .../autogen_agentchat/conditions/__init__.py | 2 + .../conditions/_terminations.py | 74 ++++++++++++++++++- .../tests/test_termination_condition.py | 56 ++++++++++++++ .../agentchat-user-guide/tutorial/teams.ipynb | 1 + .../tutorial/termination.ipynb | 5 +- 5 files changed, 135 insertions(+), 3 deletions(-) diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py index 0a53f73984..72b61745ac 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py +++ b/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/__init__.py @@ -5,6 +5,7 @@ multi-agent teams. from ._terminations import ( ExternalTermination, + FunctionalTermination, FunctionCallTermination, HandoffTermination, MaxMessageTermination, @@ -27,4 +28,5 @@ __all__ = [ "SourceMatchTermination", "TextMessageTermination", "FunctionCallTermination", + "FunctionalTermination", ] 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 82ff975b67..f0ba274ebe 100644 --- a/python/packages/autogen-agentchat/src/autogen_agentchat/conditions/_terminations.py +++ b/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 @@ class TextMentionTermination(TerminationCondition, Component[TextMentionTerminat 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 diff --git a/python/packages/autogen-agentchat/tests/test_termination_condition.py b/python/packages/autogen-agentchat/tests/test_termination_condition.py index 230a8d4ac7..bccb2f0124 100644 --- a/python/packages/autogen-agentchat/tests/test_termination_condition.py +++ b/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 @@ from autogen_agentchat.conditions import ( 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 diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb index e5419c0a89..2ae406ea14 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/teams.ipynb +++ b/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", diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb index 895303fe44..b12b874043 100644 --- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/termination.ipynb +++ b/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,