Migrate team one to new subscriptions (#560)

Co-authored-by: afourney <adam.fourney@gmail.com>

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

@@ -3,7 +3,7 @@ The :mod:`autogen_core.components` module provides building blocks for creating
 """
 
 from ._closure_agent import ClosureAgent
-from ._default_subscription import DefaultSubscription, default_subscription
+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
@@ -23,4 +23,5 @@ __all__ = [
     "DefaultSubscription",
     "DefaultTopicId",
     "default_subscription",
+    "type_subscription",
 ]

python/packages/team-one/examples/example.py

@@ -10,7 +10,6 @@ import logging
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.base import AgentId, AgentProxy
-from autogen_core.components import DefaultSubscription
 from autogen_core.components.code_executor._impl.docker_command_line_code_executor import DockerCommandLineCodeExecutor
 from team_one.agents.coder import Coder, Executor
 from team_one.agents.orchestrator import LedgerOrchestrator
@@ -25,32 +24,30 @@ async def main() -> None:
 
     async with DockerCommandLineCodeExecutor() as code_executor:
         # Register agents.
-        await runtime.register(
-            "Coder", lambda: Coder(model_client=create_completion_client_from_env()), lambda: [DefaultSubscription()]
-        )
+        await Coder.register(runtime, "Coder", lambda: Coder(model_client=create_completion_client_from_env()))
         coder = AgentProxy(AgentId("Coder", "default"), runtime)
 
-        await runtime.register(
+        await Executor.register(
+            runtime,
             "Executor",
             lambda: Executor("A agent for executing code", executor=code_executor),
-            lambda: [DefaultSubscription()],
         )
         executor = AgentProxy(AgentId("Executor", "default"), runtime)
 
-        await runtime.register(
+        await UserProxy.register(
+            runtime,
             "UserProxy",
             lambda: UserProxy(description="The current user interacting with you."),
-            lambda: [DefaultSubscription()],
         )
         user_proxy = AgentProxy(AgentId("UserProxy", "default"), runtime)
 
         # TODO: doesn't work for more than default key
-        await runtime.register(
+        await LedgerOrchestrator.register(
+            runtime,
             "orchestrator",
             lambda: LedgerOrchestrator(
                 model_client=create_completion_client_from_env(), agents=[coder, executor, user_proxy]
             ),
-            lambda: [DefaultSubscription()],
         )
 
         runtime.start()

python/packages/team-one/examples/example_coder.py

@@ -10,7 +10,6 @@ import logging
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.base import AgentId, AgentProxy
-from autogen_core.components import DefaultSubscription
 from autogen_core.components.code_executor._impl.docker_command_line_code_executor import DockerCommandLineCodeExecutor
 from team_one.agents.coder import Coder, Executor
 from team_one.agents.orchestrator import RoundRobinOrchestrator
@@ -25,29 +24,27 @@ async def main() -> None:
 
     async with DockerCommandLineCodeExecutor() as code_executor:
         # Register agents.
-        await runtime.register(
-            "Coder", lambda: Coder(model_client=create_completion_client_from_env()), lambda: [DefaultSubscription()]
-        )
+        await Coder.register(runtime, "Coder", lambda: Coder(model_client=create_completion_client_from_env()))
         coder = AgentProxy(AgentId("Coder", "default"), runtime)
 
-        await runtime.register(
+        await Executor.register(
+            runtime,
             "Executor",
             lambda: Executor("A agent for executing code", executor=code_executor),
-            lambda: [DefaultSubscription()],
         )
         executor = AgentProxy(AgentId("Executor", "default"), runtime)
 
-        await runtime.register(
+        await UserProxy.register(
+            runtime,
             "UserProxy",
             lambda: UserProxy(description="The current user interacting with you."),
-            lambda: [DefaultSubscription()],
         )
         user_proxy = AgentProxy(AgentId("UserProxy", "default"), runtime)
 
-        await runtime.register(
+        await RoundRobinOrchestrator.register(
+            runtime,
             "orchestrator",
             lambda: RoundRobinOrchestrator([coder, executor, user_proxy]),
-            lambda: [DefaultSubscription()],
         )
 
         runtime.start()

python/packages/team-one/examples/example_file_surfer.py

@@ -8,7 +8,6 @@ import logging
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.base import AgentId, AgentProxy
-from autogen_core.components import DefaultSubscription
 from team_one.agents.file_surfer import FileSurfer
 from team_one.agents.orchestrator import RoundRobinOrchestrator
 from team_one.agents.user_proxy import UserProxy
@@ -24,14 +23,14 @@ async def main() -> None:
     client = create_completion_client_from_env()
 
     # Register agents.
-    await runtime.register("file_surfer", lambda: FileSurfer(model_client=client), lambda: [DefaultSubscription()])
+    await FileSurfer.register(runtime, "file_surfer", lambda: FileSurfer(model_client=client))
     file_surfer = AgentProxy(AgentId("file_surfer", "default"), runtime)
 
-    await runtime.register("UserProxy", lambda: UserProxy(), lambda: [DefaultSubscription()])
+    await UserProxy.register(runtime, "UserProxy", lambda: UserProxy())
     user_proxy = AgentProxy(AgentId("UserProxy", "default"), runtime)
 
-    await runtime.register(
-        "orchestrator", lambda: RoundRobinOrchestrator([file_surfer, user_proxy]), lambda: [DefaultSubscription()]
+    await RoundRobinOrchestrator.register(
+        runtime, "orchestrator", lambda: RoundRobinOrchestrator([file_surfer, user_proxy])
     )
 
     runtime.start()

python/packages/team-one/examples/example_reflexagents.py

@@ -7,7 +7,7 @@ import logging
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.base import AgentId, AgentProxy
-from autogen_core.components import DefaultSubscription, DefaultTopicId
+from autogen_core.components import DefaultTopicId
 from autogen_core.components.models import UserMessage
 from team_one.agents.orchestrator import RoundRobinOrchestrator
 from team_one.agents.reflex_agents import ReflexAgent
@@ -18,17 +18,17 @@ from team_one.utils import LogHandler
 async def main() -> None:
     runtime = SingleThreadedAgentRuntime()
 
-    await runtime.register("fake_agent_1", lambda: ReflexAgent("First reflect agent"), lambda: [DefaultSubscription()])
+    await ReflexAgent.register(runtime, "fake_agent_1", lambda: ReflexAgent("First reflect agent"))
     fake1 = AgentProxy(AgentId("fake_agent_1", "default"), runtime)
 
-    await runtime.register("fake_agent_2", lambda: ReflexAgent("Second reflect agent"), lambda: [DefaultSubscription()])
+    await ReflexAgent.register(runtime, "fake_agent_2", lambda: ReflexAgent("Second reflect agent"))
     fake2 = AgentProxy(AgentId("fake_agent_2", "default"), runtime)
 
-    await runtime.register("fake_agent_3", lambda: ReflexAgent("Third reflect agent"), lambda: [DefaultSubscription()])
+    await ReflexAgent.register(runtime, "fake_agent_3", lambda: ReflexAgent("Third reflect agent"))
     fake3 = AgentProxy(AgentId("fake_agent_3", "default"), runtime)
 
-    await runtime.register(
-        "orchestrator", lambda: RoundRobinOrchestrator([fake1, fake2, fake3]), lambda: [DefaultSubscription()]
+    await RoundRobinOrchestrator.register(
+        runtime, "orchestrator", lambda: RoundRobinOrchestrator([fake1, fake2, fake3])
     )
 
     task_message = UserMessage(content="Test Message", source="User")

python/packages/team-one/examples/example_userproxy.py

@@ -12,7 +12,6 @@ from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.base import AgentId, AgentProxy
 
 # from typing import Any, Dict, List, Tuple, Union
-from autogen_core.components import DefaultSubscription
 from team_one.agents.coder import Coder
 from team_one.agents.orchestrator import RoundRobinOrchestrator
 from team_one.agents.user_proxy import UserProxy
@@ -28,15 +27,13 @@ async def main() -> None:
     client = create_completion_client_from_env()
 
     # Register agents.
-    await runtime.register("Coder", lambda: Coder(model_client=client), lambda: [DefaultSubscription()])
+    await Coder.register(runtime, "Coder", lambda: Coder(model_client=client))
     coder = AgentProxy(AgentId("Coder", "default"), runtime)
 
-    await runtime.register("UserProxy", lambda: UserProxy(), lambda: [DefaultSubscription()])
+    await UserProxy.register(runtime, "UserProxy", lambda: UserProxy())
     user_proxy = AgentProxy(AgentId("UserProxy", "default"), runtime)
 
-    await runtime.register(
-        "orchestrator", lambda: RoundRobinOrchestrator([coder, user_proxy]), lambda: [DefaultSubscription()]
-    )
+    await RoundRobinOrchestrator.register(runtime, "orchestrator", lambda: RoundRobinOrchestrator([coder, user_proxy]))
 
     runtime.start()
     await runtime.send_message(RequestReplyMessage(), user_proxy.id)

python/packages/team-one/examples/example_websurfer.py

@@ -10,7 +10,6 @@ import os
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.base import AgentId, AgentProxy
-from autogen_core.components import DefaultSubscription
 from team_one.agents.multimodal_web_surfer import MultimodalWebSurfer
 from team_one.agents.orchestrator import RoundRobinOrchestrator
 from team_one.agents.user_proxy import UserProxy
@@ -28,14 +27,14 @@ async def main() -> None:
     client = create_completion_client_from_env(model="gpt-4o")
 
     # Register agents.
-    await runtime.register("WebSurfer", lambda: MultimodalWebSurfer(), lambda: [DefaultSubscription()])
+    await MultimodalWebSurfer.register(runtime, "WebSurfer", MultimodalWebSurfer)
     web_surfer = AgentProxy(AgentId("WebSurfer", "default"), runtime)
 
-    await runtime.register("UserProxy", lambda: UserProxy(), lambda: [DefaultSubscription()])
+    await UserProxy.register(runtime, "UserProxy", UserProxy)
     user_proxy = AgentProxy(AgentId("UserProxy", "default"), runtime)
 
-    await runtime.register(
-        "orchestrator", lambda: RoundRobinOrchestrator([web_surfer, user_proxy]), lambda: [DefaultSubscription()]
+    await RoundRobinOrchestrator.register(
+        runtime, "orchestrator", lambda: RoundRobinOrchestrator([web_surfer, user_proxy])
     )
 
     runtime.start()

python/packages/team-one/pyproject.toml

@@ -74,5 +74,5 @@ include = [
 
 
 [tool.pyright]
-extend = "../../pyproject.toml"
+extends = "../../pyproject.toml"
 include = ["src", "tests", "examples"]

python/packages/team-one/src/team_one/agents/coder.py

@@ -1,8 +1,9 @@
 import re
-from typing import List, Optional, Tuple, Union
+from typing import List, Tuple, Union
 
 from autogen_core.base import CancellationToken
-from autogen_core.components.code_executor import CodeBlock, CodeExecutor, DockerCommandLineCodeExecutor
+from autogen_core.components import default_subscription
+from autogen_core.components.code_executor import CodeBlock, CodeExecutor
 from autogen_core.components.models import (
     ChatCompletionClient,
     SystemMessage,
@@ -14,6 +15,7 @@ from ..utils import message_content_to_str
 from .base_worker import BaseWorker
 
 
+@default_subscription
 class Coder(BaseWorker):
     """An agent that uses tools to write, execute, and debug Python code."""
 
@@ -52,6 +54,7 @@ Reply "TERMINATE" in the end when everything is done.""")
         return "TERMINATE" in response.content, response.content
 
 
+@default_subscription
 class Executor(BaseWorker):
     DEFAULT_DESCRIPTION = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks)"
 

python/packages/team-one/src/team_one/agents/file_surfer/file_surfer.py

@@ -3,7 +3,7 @@ import time
 from typing import List, Optional, Tuple
 
 from autogen_core.base import CancellationToken
-from autogen_core.components import FunctionCall
+from autogen_core.components import FunctionCall, default_subscription
 from autogen_core.components.models import (
     ChatCompletionClient,
     SystemMessage,
@@ -17,6 +17,7 @@ from ..base_worker import BaseWorker
 from ._tools import TOOL_FIND_NEXT, TOOL_FIND_ON_PAGE_CTRL_F, TOOL_OPEN_LOCAL_FILE, TOOL_PAGE_DOWN, TOOL_PAGE_UP
 
 
+@default_subscription
 class FileSurfer(BaseWorker):
     """An agent that uses tools to read and navigate local files."""
 

python/packages/team-one/src/team_one/agents/multimodal_web_surfer/multimodal_web_surfer.py

@@ -13,7 +13,7 @@ from urllib.parse import quote_plus  # parse_qs, quote, unquote, urlparse, urlun
 import aiofiles
 from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.base import CancellationToken
-from autogen_core.components import FunctionCall
+from autogen_core.components import FunctionCall, default_subscription
 from autogen_core.components import Image as AGImage
 from autogen_core.components.models import (
     AssistantMessage,
@@ -75,6 +75,7 @@ class DEFAULT_CHANNEL(metaclass=SentinelMeta):
     pass
 
 
+@default_subscription
 class MultimodalWebSurfer(BaseWorker):
     """(In preview) A multimodal agent that acts as a web surfer that can search the web and visit web pages."""
 

python/packages/team-one/src/team_one/agents/orchestrator.py

@@ -2,6 +2,7 @@ import json
 from typing import Any, Dict, List, Optional
 
 from autogen_core.base import AgentProxy, MessageContext, TopicId
+from autogen_core.components import default_subscription
 from autogen_core.components.models import (
     AssistantMessage,
     ChatCompletionClient,
@@ -23,6 +24,7 @@ from .orchestrator_prompts import (
 )
 
 
+@default_subscription
 class RoundRobinOrchestrator(BaseOrchestrator):
     def __init__(
         self,
@@ -37,6 +39,7 @@ class RoundRobinOrchestrator(BaseOrchestrator):
         return self._agents[self._current_index]
 
 
+@default_subscription
 class LedgerOrchestrator(BaseOrchestrator):
     DEFAULT_SYSTEM_MESSAGES = [
         SystemMessage(ORCHESTRATOR_SYSTEM_MESSAGE),

python/packages/team-one/src/team_one/agents/reflex_agents.py

@@ -1,10 +1,11 @@
 from autogen_core.base import MessageContext, TopicId
-from autogen_core.components import RoutedAgent, message_handler
+from autogen_core.components import RoutedAgent, default_subscription, message_handler
 from autogen_core.components.models import UserMessage
 
 from ..messages import BroadcastMessage, RequestReplyMessage
 
 
+@default_subscription
 class ReflexAgent(RoutedAgent):
     def __init__(self, description: str) -> None:
         super().__init__(description)

python/packages/team-one/src/team_one/agents/user_proxy.py

@@ -2,11 +2,13 @@ import asyncio
 from typing import Tuple
 
 from autogen_core.base import CancellationToken
+from autogen_core.components import default_subscription
 
 from ..messages import UserContent
 from .base_worker import BaseWorker
 
 
+@default_subscription
 class UserProxy(BaseWorker):
     """An agent that allows the user to play the role of an agent in the conversation."""
 

python/packages/team-one/tests/headless_web_surfer/test_web_surfer.py

@@ -96,7 +96,8 @@ async def test_web_surfer() -> None:
     # Register agents.
 
     # Register agents.
-    await runtime.register(
+    await MultimodalWebSurfer.register(
+        runtime,
         "WebSurfer",
         lambda: MultimodalWebSurfer(),
     )
@@ -168,18 +169,22 @@ async def test_web_surfer_oai() -> None:
     client = create_completion_client_from_env()
 
     # Register agents.
-    await runtime.register(
+    await MultimodalWebSurfer.register(
+        runtime,
         "WebSurfer",
         lambda: MultimodalWebSurfer(),
     )
     web_surfer = AgentProxy(AgentId("WebSurfer", "default"), runtime)
 
-    await runtime.register(
+    await UserProxy.register(
+        runtime,
         "UserProxy",
         lambda: UserProxy(),
     )
     user_proxy = AgentProxy(AgentId("UserProxy", "default"), runtime)
-    await runtime.register("orchestrator", lambda: RoundRobinOrchestrator([web_surfer, user_proxy]))
+    await RoundRobinOrchestrator.register(
+        runtime, "orchestrator", lambda: RoundRobinOrchestrator([web_surfer, user_proxy])
+    )
     runtime.start()
 
     actual_surfer = await runtime.try_get_underlying_agent_instance(web_surfer.id, MultimodalWebSurfer)
@@ -233,9 +238,8 @@ async def test_web_surfer_bing() -> None:
     client = create_completion_client_from_env(env)
 
     # Register agents.
-
-    # Register agents.
-    await runtime.register(
+    await MultimodalWebSurfer.register(
+        runtime,
         "WebSurfer",
         lambda: MultimodalWebSurfer(),
     )