Skip to content

Commit

Permalink
Change base agent to use on message impl (#4485)
Browse files Browse the repository at this point in the history
* Change base agent to use on message impl

* update doc

* Update

* Merge branch 'main' into on_msg_impl
  • Loading branch information
jackgerrits authored Dec 3, 2024
1 parent 934ae03 commit 954ba84
Show file tree
Hide file tree
Showing 8 changed files with 21 additions and 17 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ def __init__(self, description: str) -> None:
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()
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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": [
Expand All @@ -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"
]
},
Expand Down Expand Up @@ -254,7 +254,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "autogen_core",
"display_name": ".venv",
"language": "python",
"name": "python3"
},
Expand All @@ -268,7 +268,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
"version": "3.12.6"
}
},
"nbformat": 4,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
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

Expand Down Expand Up @@ -108,8 +108,12 @@ def id(self) -> AgentId:
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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ def id(self) -> AgentId:
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}"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -470,7 +470,7 @@ def __init__(self, description: str) -> None:

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."""
Expand Down
2 changes: 1 addition & 1 deletion python/packages/autogen-core/tests/test_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ def __init__(self) -> None:
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]:
Expand Down
6 changes: 3 additions & 3 deletions python/packages/autogen-core/tests/test_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
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


Expand All @@ -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
Expand All @@ -37,7 +37,7 @@ async def handler2(self, message: str | bool, ctx: MessageContext) -> 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
Expand Down
2 changes: 1 addition & 1 deletion python/packages/autogen-core/tests/test_utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 954ba84

Please sign in to comment.