diff --git a/python/packages/autogen-core/src/autogen_core/application/_worker_runtime_host_servicer.py b/python/packages/autogen-core/src/autogen_core/application/_worker_runtime_host_servicer.py index 5cd2bf8ea9b8..7c597bd07a8f 100644 --- a/python/packages/autogen-core/src/autogen_core/application/_worker_runtime_host_servicer.py +++ b/python/packages/autogen-core/src/autogen_core/application/_worker_runtime_host_servicer.py @@ -4,7 +4,7 @@ from asyncio import Future, Task from typing import Any, Dict, Set -from autogen_core.components._type_prefix_subscription import TypePrefixSubscription +from autogen_core.base._type_prefix_subscription import TypePrefixSubscription from ..base import Subscription, TopicId from ..components import TypeSubscription diff --git a/python/packages/autogen-core/src/autogen_core/base/_base_agent.py b/python/packages/autogen-core/src/autogen_core/base/_base_agent.py index 5d8e94225b3a..70481705ca6e 100644 --- a/python/packages/autogen-core/src/autogen_core/base/_base_agent.py +++ b/python/packages/autogen-core/src/autogen_core/base/_base_agent.py @@ -20,6 +20,7 @@ from ._subscription import Subscription, UnboundSubscription from ._subscription_context import SubscriptionInstantiationContext from ._topic import TopicId +from ._type_prefix_subscription import TypePrefixSubscription T = TypeVar("T", bound=Agent) @@ -149,6 +150,7 @@ async def register( factory: Callable[[], Self | Awaitable[Self]], *, skip_class_subscriptions: bool = False, + skip_direct_message_subscription: bool = False, ) -> AgentType: agent_type = AgentType(type) agent_type = await runtime.register_factory(type=agent_type, agent_factory=factory, expected_class=cls) @@ -166,6 +168,16 @@ async def register( for subscription in subscriptions: await runtime.add_subscription(subscription) + if not skip_direct_message_subscription: + # Additionally adds a special prefix subscription for this agent to receive direct messages + await runtime.add_subscription( + TypePrefixSubscription( + # The prefix MUST include ":" to avoid collisions with other agents + topic_type_prefix=agent_type.type + ":", + agent_type=agent_type.type, + ) + ) + # TODO: deduplication for _message_type, serializer in cls._handles_types(): runtime.add_message_serializer(serializer) diff --git a/python/packages/autogen-core/src/autogen_core/components/_type_prefix_subscription.py b/python/packages/autogen-core/src/autogen_core/base/_type_prefix_subscription.py similarity index 93% rename from python/packages/autogen-core/src/autogen_core/components/_type_prefix_subscription.py rename to python/packages/autogen-core/src/autogen_core/base/_type_prefix_subscription.py index d71b587d1418..f00119165090 100644 --- a/python/packages/autogen-core/src/autogen_core/components/_type_prefix_subscription.py +++ b/python/packages/autogen-core/src/autogen_core/base/_type_prefix_subscription.py @@ -1,7 +1,9 @@ import uuid -from ..base import AgentId, Subscription, TopicId -from ..base.exceptions import CantHandleException +from ._agent_id import AgentId +from ._subscription import Subscription +from ._topic import TopicId +from .exceptions import CantHandleException class TypePrefixSubscription(Subscription): diff --git a/python/packages/autogen-core/src/autogen_core/components/__init__.py b/python/packages/autogen-core/src/autogen_core/components/__init__.py index 4c4d02f3be25..75bb5eabcbe8 100644 --- a/python/packages/autogen-core/src/autogen_core/components/__init__.py +++ b/python/packages/autogen-core/src/autogen_core/components/__init__.py @@ -2,12 +2,12 @@ The :mod:`autogen_core.components` module provides building blocks for creating single agents """ +from ..base._type_prefix_subscription import TypePrefixSubscription from ._closure_agent import ClosureAgent from ._default_subscription import DefaultSubscription, default_subscription, type_subscription from ._default_topic import DefaultTopicId from ._image import Image from ._routed_agent import RoutedAgent, TypeRoutedAgent, event, message_handler, rpc -from ._type_prefix_subscription import TypePrefixSubscription from ._type_subscription import TypeSubscription from ._types import FunctionCall diff --git a/python/packages/autogen-core/tests/test_worker_runtime.py b/python/packages/autogen-core/tests/test_worker_runtime.py index d58233c3ac82..26c95dc01860 100644 --- a/python/packages/autogen-core/tests/test_worker_runtime.py +++ b/python/packages/autogen-core/tests/test_worker_runtime.py @@ -360,7 +360,7 @@ async def get_subscribed_recipients() -> List[AgentId]: ) subscriptions1 = get_current_subscriptions() - assert len(subscriptions1) == 1 + assert len(subscriptions1) == 2 recipients1 = await get_subscribed_recipients() assert AgentId(type="worker1", key="default") in recipients1 @@ -388,7 +388,7 @@ async def get_subscribed_recipients() -> List[AgentId]: ) subscriptions3 = get_current_subscriptions() - assert len(subscriptions3) == 1 + assert len(subscriptions3) == 2 assert first_subscription_id not in [x.id for x in subscriptions3] recipients3 = await get_subscribed_recipients()