Make termination condition config declarative

python/packages/autogen-agentchat/src/autogen_agentchat/base/_termination.py

@@ -1,14 +1,19 @@
 import asyncio
 from abc import ABC, abstractmethod
-from typing import List, Sequence
+from typing import Any, List, Sequence
+from typing_extensions import Self
+
+from pydantic import BaseModel
 
 from ..messages import AgentEvent, ChatMessage, StopMessage
+from autogen_core import Component, ComponentModel, ComponentLoader
 
 
-class TerminatedException(BaseException): ...
+class TerminatedException(BaseException):
+    ...
 
 
-class TerminationCondition(ABC):
+class TerminationCondition(ABC, Component[Any]):
     """A stateful condition that determines when a conversation should be terminated.
 
     A termination condition is a callable that takes a sequence of ChatMessage objects
@@ -43,6 +48,8 @@ async def main() -> None:
             asyncio.run(main())
     """
 
+    component_type = "termination"
+
     @property
     @abstractmethod
     def terminated(self) -> bool:
@@ -79,7 +86,15 @@ def __or__(self, other: "TerminationCondition") -> "TerminationCondition":
         return _OrTerminationCondition(self, other)
 
 
-class _AndTerminationCondition(TerminationCondition):
+class AndTerminationConditionConfig(BaseModel):
+    conditions:  List[ComponentModel]
+
+
+class _AndTerminationCondition(TerminationCondition, Component[AndTerminationConditionConfig]):
+
+    component_config_schema = AndTerminationConditionConfig
+    component_type = "termination"
+
     def __init__(self, *conditions: TerminationCondition) -> None:
         self._conditions = conditions
         self._stop_messages: List[StopMessage] = []
@@ -90,7 +105,8 @@ def terminated(self) -> bool:
 
     async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None:
         if self.terminated:
-            raise TerminatedException("Termination condition has already been reached.")
+            raise TerminatedException(
+                "Termination condition has already been reached.")
         # Check all remaining conditions.
         stop_messages = await asyncio.gather(
             *[condition(messages) for condition in self._conditions if not condition.terminated]
@@ -102,17 +118,44 @@ async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMe
         if any(stop_message is None for stop_message in stop_messages):
             # If any remaining condition has not reached termination, it is not terminated.
             return None
-        content = ", ".join(stop_message.content for stop_message in self._stop_messages)
-        source = ", ".join(stop_message.source for stop_message in self._stop_messages)
+        content = ", ".join(
+            stop_message.content for stop_message in self._stop_messages)
+        source = ", ".join(
+            stop_message.source for stop_message in self._stop_messages)
         return StopMessage(content=content, source=source)
 
     async def reset(self) -> None:
         for condition in self._conditions:
             await condition.reset()
         self._stop_messages.clear()
 
+    def _to_config(self) -> AndTerminationConditionConfig:
+        """Convert the AND termination condition to a config."""
+        return AndTerminationConditionConfig(
+            conditions=[condition.dump_component()
+                        for condition in self._conditions]
+        )
+
+    @classmethod
+    def _from_config(cls, config: AndTerminationConditionConfig) -> Self:
+        """Create an AND termination condition from a config."""
+        conditions = [
+            ComponentLoader.load_component(
+                condition_model, TerminationCondition)
+            for condition_model in config.conditions
+        ]
+        return cls(*conditions)
+
+
+class OrTerminationConditionConfig(BaseModel):
+    conditions: List[ComponentModel]
+    """List of termination conditions where any one being satisfied is sufficient."""
+
+
+class _OrTerminationCondition(TerminationCondition, Component[OrTerminationConditionConfig]):
+    component_config_schema = OrTerminationConditionConfig
+    component_type = "termination"
 
-class _OrTerminationCondition(TerminationCondition):
     def __init__(self, *conditions: TerminationCondition) -> None:
         self._conditions = conditions
 
@@ -122,14 +165,34 @@ def terminated(self) -> bool:
 
     async def __call__(self, messages: Sequence[AgentEvent | ChatMessage]) -> StopMessage | None:
         if self.terminated:
-            raise RuntimeError("Termination condition has already been reached")
+            raise RuntimeError(
+                "Termination condition has already been reached")
         stop_messages = await asyncio.gather(*[condition(messages) for condition in self._conditions])
         if any(stop_message is not None for stop_message in stop_messages):
-            content = ", ".join(stop_message.content for stop_message in stop_messages if stop_message is not None)
-            source = ", ".join(stop_message.source for stop_message in stop_messages if stop_message is not None)
+            content = ", ".join(
+                stop_message.content for stop_message in stop_messages if stop_message is not None)
+            source = ", ".join(
+                stop_message.source for stop_message in stop_messages if stop_message is not None)
             return StopMessage(content=content, source=source)
         return None
 
     async def reset(self) -> None:
         for condition in self._conditions:
             await condition.reset()
+
+    def _to_config(self) -> OrTerminationConditionConfig:
+        """Convert the OR termination condition to a config."""
+        return OrTerminationConditionConfig(
+            conditions=[condition.dump_component()
+                        for condition in self._conditions]
+        )
+
+    @classmethod
+    def _from_config(cls, config: OrTerminationConditionConfig) -> Self:
+        """Create an OR termination condition from a config."""
+        conditions = [
+            ComponentLoader.load_component(
+                condition_model, TerminationCondition)
+            for condition_model in config.conditions
+        ]
+        return cls(*conditions)