Change base agent to use on message impl (#4485)

* Change base agent to use on message impl

* update doc

* Update

* Merge branch 'main' into on_msg_impl

python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_sequential_routed_agent.py

@@ -43,9 +43,9 @@ class SequentialRoutedAgent(RoutedAgent):
         super().__init__(description=description)
         self._fifo_lock = FIFOLock()
 
-    async def on_message(self, message: Any, ctx: MessageContext) -> Any | None:
+    async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any | None:
         await self._fifo_lock.acquire()
         try:
-            return await super().on_message(message, ctx)
+            return await super().on_message_impl(message, ctx)
         finally:
             self._fifo_lock.release()

python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/agent-and-agent-runtime.ipynb

@@ -19,7 +19,7 @@
     "An agent in AutoGen is an entity defined by the base class {py:class}`autogen_core.base.BaseAgent`.\n",
     "It has a unique identifier of the type {py:class}`autogen_core.base.AgentId`,\n",
     "a metadata dictionary of the type {py:class}`autogen_core.base.AgentMetadata`,\n",
-    "and method for handling messages {py:meth}`autogen_core.base.BaseAgent.on_message`.\n",
+    "and method for handling messages {py:meth}`autogen_core.base.BaseAgent.on_message_impl`.\n",
     "\n",
     "An agent runtime is the execution environment for agents in AutoGen.\n",
     "Similar to the runtime environment of a programming language,\n",
@@ -42,14 +42,14 @@
     "## Implementing an Agent\n",
     "\n",
     "To implement an agent, the developer must subclass the {py:class}`~autogen_core.base.BaseAgent` class\n",
-    "and implement the {py:meth}`~autogen_core.base.BaseAgent.on_message` method.\n",
+    "and implement the {py:meth}`~autogen_core.base.BaseAgent.on_message_impl` method.\n",
     "This method is invoked when the agent receives a message. For example,\n",
     "the following agent handles a simple message type and prints the message it receives:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -67,7 +67,7 @@
     "    def __init__(self) -> None:\n",
     "        super().__init__(\"MyAgent\")\n",
     "\n",
-    "    async def on_message(self, message: MyMessageType, ctx: MessageContext) -> None:\n",
+    "    async def on_message_impl(self, message: MyMessageType, ctx: MessageContext) -> None:\n",
     "        print(f\"Received message: {message.content}\")  # type: ignore"
    ]
   },
@@ -254,7 +254,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "autogen_core",
+   "display_name": ".venv",
    "language": "python",
    "name": "python3"
   },
@@ -268,7 +268,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.9"
+   "version": "3.12.6"
   }
  },
  "nbformat": 4,

python/packages/autogen-core/src/autogen_core/base/_base_agent.py

@@ -4,7 +4,7 @@ import inspect
 import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from typing import Any, Awaitable, Callable, ClassVar, List, Mapping, Tuple, Type, TypeVar
+from typing import Any, Awaitable, Callable, ClassVar, List, Mapping, Tuple, Type, TypeVar, final
 
 from typing_extensions import Self
 
@@ -108,8 +108,12 @@ class BaseAgent(ABC, Agent):
     def runtime(self) -> AgentRuntime:
         return self._runtime
 
+    @final
+    async def on_message(self, message: Any, ctx: MessageContext) -> Any:
+        return await self.on_message_impl(message, ctx)
+
     @abstractmethod
-    async def on_message(self, message: Any, ctx: MessageContext) -> Any: ...
+    async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any: ...
 
     async def send_message(
         self,

python/packages/autogen-core/src/autogen_core/components/_closure_agent.py

@@ -111,7 +111,7 @@ class ClosureAgent(BaseAgent, ClosureContext):
     def runtime(self) -> AgentRuntime:
         return self._runtime
 
-    async def on_message(self, message: Any, ctx: MessageContext) -> Any:
+    async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any:
         if type(message) not in self._expected_types:
             raise CantHandleException(
                 f"Message type {type(message)} not in target types {self._expected_types} of {self.id}"

python/packages/autogen-core/src/autogen_core/components/_routed_agent.py

@@ -470,7 +470,7 @@ class RoutedAgent(BaseAgent):
 
         super().__init__(description)
 
-    async def on_message(self, message: Any, ctx: MessageContext) -> Any | None:
+    async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any | None:
         """Handle a message by routing it to the appropriate message handler.
         Do not override this method in subclasses. Instead, add message handlers as methods decorated with
         either the :func:`event` or :func:`rpc` decorator."""

python/packages/autogen-core/tests/test_state.py

@@ -10,7 +10,7 @@ class StatefulAgent(BaseAgent):
         super().__init__("A stateful agent")
         self.state = 0
 
-    async def on_message(self, message: Any, ctx: MessageContext) -> None:
+    async def on_message_impl(self, message: Any, ctx: MessageContext) -> None:
         raise NotImplementedError
 
     async def save_state(self) -> Mapping[str, Any]:

python/packages/autogen-core/tests/test_types.py

@@ -5,7 +5,7 @@ from typing import Any, List, Optional, Union
 from autogen_core.base import MessageContext
 from autogen_core.base._serialization import has_nested_base_model
 from autogen_core.base._type_helpers import AnyType, get_types
-from autogen_core.components._routed_agent import message_handler
+from autogen_core.components._routed_agent import RoutedAgent, message_handler
 from pydantic import BaseModel
 
 
@@ -21,7 +21,7 @@ def test_get_types() -> None:
 
 
 def test_handler() -> None:
-    class HandlerClass:
+    class HandlerClass(RoutedAgent):
         @message_handler()
         async def handler(self, message: int, ctx: MessageContext) -> Any:
             return None
@@ -37,7 +37,7 @@ def test_handler() -> None:
     assert HandlerClass.handler2.produces_types == [NoneType]
 
 
-class HandlerClass:
+class HandlerClass(RoutedAgent):
     @message_handler()
     async def handler(self, message: int, ctx: MessageContext) -> Any:
         return None

python/packages/autogen-core/tests/test_utils/__init__.py

@@ -57,5 +57,5 @@ class NoopAgent(BaseAgent):
     def __init__(self) -> None:
         super().__init__("A no op agent")
 
-    async def on_message(self, message: Any, ctx: MessageContext) -> Any:
+    async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any:
         raise NotImplementedError