Flatten core base and components

python/packages/agbench/benchmarks/AssistantBench/Templates/MagenticOne/scenario.py

@@ -8,7 +8,7 @@ from openai import AzureOpenAI
 
 from typing import List
 
-from autogen_core.base import AgentId, AgentProxy, TopicId
+from autogen_core import AgentId, AgentProxy, TopicId
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.components.models import (
@@ -16,7 +16,7 @@ from autogen_core.components.models import (
     UserMessage,
     LLMMessage,
 )
-from autogen_core.components import DefaultSubscription, DefaultTopicId
+from autogen_core import DefaultSubscription, DefaultTopicId
 from autogen_core.components.code_executor import LocalCommandLineCodeExecutor
 from autogen_core.components.models import AssistantMessage
 

python/packages/agbench/benchmarks/GAIA/Templates/MagenticOne/scenario.py

@@ -8,7 +8,7 @@ from openai import AzureOpenAI
 
 from typing import List
 
-from autogen_core.base import AgentId, AgentProxy, TopicId
+from autogen_core import AgentId, AgentProxy, TopicId
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
 from autogen_core.components.models import (
@@ -17,7 +17,7 @@ from autogen_core.components.models import (
     UserMessage,
     LLMMessage,
 )
-from autogen_core.components import DefaultSubscription, DefaultTopicId
+from autogen_core import DefaultSubscription, DefaultTopicId
 from autogen_core.components.code_executor import LocalCommandLineCodeExecutor
 from autogen_core.components.models import AssistantMessage
 

python/packages/agbench/benchmarks/HumanEval/Templates/MagenticOne/scenario.py

@@ -1,10 +1,10 @@
 import asyncio
 import logging
 
-from autogen_core.base import AgentId, AgentProxy, TopicId
+from autogen_core import AgentId, AgentProxy, TopicId
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
-from autogen_core.components import DefaultSubscription, DefaultTopicId
+from autogen_core import DefaultSubscription, DefaultTopicId
 from autogen_core.components.code_executor import LocalCommandLineCodeExecutor
 from autogen_core.components.models import (
     UserMessage,
@@ -41,7 +41,7 @@ async def main() -> None:
     executor = AgentProxy(AgentId("Executor", "default"), runtime)
 
     await runtime.register(
-        "Orchestrator", 
+        "Orchestrator",
         lambda: RoundRobinOrchestrator([coder, executor]),
         subscriptions=lambda: [DefaultSubscription()],
     )

python/packages/agbench/benchmarks/WebArena/Templates/MagenticOne/scenario.py

@@ -7,10 +7,10 @@ import nltk
 
 from typing import Any, Dict, List, Tuple, Union
 
-from autogen_core.base import AgentId, AgentProxy, TopicId
+from autogen_core import AgentId, AgentProxy, TopicId
 from autogen_core.application import SingleThreadedAgentRuntime
 from autogen_core.application.logging import EVENT_LOGGER_NAME
-from autogen_core.components import DefaultSubscription, DefaultTopicId
+from autogen_core import DefaultSubscription, DefaultTopicId
 from autogen_core.components.code_executor import LocalCommandLineCodeExecutor
 from autogen_core.components.models import (
     ChatCompletionClient,
@@ -105,7 +105,7 @@ async def main() -> None:
             task_prompt = task_prompt.replace(k, REPLACEMENTS[k])
         fh.write(task_prompt)
         TASK = json.loads(task_prompt)
-        if TASK["start_url"] == REDDIT: 
+        if TASK["start_url"] == REDDIT:
             TASK["start_url"] = TASK["start_url"] + "/forums/all"
 
     full_task = ""
@@ -150,7 +150,7 @@ Once the user has taken the final necessary action to complete the task, and you
 
     # Round-robin orchestrator
     await runtime.register(
-        "round_robin_orc", 
+        "round_robin_orc",
         lambda: RoundRobinOrchestrator(agents=[web_surfer, login_assistant],),
         subscriptions=lambda: [DefaultSubscription()],
     )
@@ -163,7 +163,7 @@ Once the user has taken the final necessary action to complete the task, and you
 
             runtime.start()
             await runtime.publish_message(
-                ResetMessage(), 
+                ResetMessage(),
                 topic_id=DefaultTopicId(),
             )
             await runtime.publish_message(
@@ -192,16 +192,16 @@ Once the user has taken the final necessary action to complete the task, and you
         subscriptions=lambda: [DefaultSubscription()],
     )
     executor = AgentProxy(AgentId("ComputerTerminal", "default"), runtime)
-    
+
     await runtime.register(
         "FileSurfer",
         lambda: FileSurfer(model_client=client),
         subscriptions=lambda: [DefaultSubscription()],
     )
     file_surfer = AgentProxy(AgentId("FileSurfer", "default"), runtime)
-   
+
     await runtime.register(
-        "orchestrator", 
+        "orchestrator",
         lambda: LedgerOrchestrator(
             agents=[coder, executor, file_surfer, web_surfer],
             model_client=client,
@@ -251,7 +251,7 @@ Once the user has taken the final necessary action to complete the task, and you
     page = actual_surfer._page
     cdp_session = await context.new_cdp_session(page)
     config_file = "full_task.json"
-    
+
     evaluator = evaluation_harness.evaluator_router(config_file)
     score = await evaluator(
         trajectory=evaluation_harness.make_answer_trajecotry(final_answer),
@@ -260,7 +260,7 @@ Once the user has taken the final necessary action to complete the task, and you
         client=cdp_session,
     #    azure_config=llm_config,
     )
-    
+
     print("FINAL SCORE: " + str(score))
 
 

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py

@@ -4,8 +4,7 @@ import logging
 import warnings
 from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Sequence
 
-from autogen_core.base import CancellationToken
-from autogen_core.components import FunctionCall
+from autogen_core import CancellationToken, FunctionCall
 from autogen_core.components.models import (
     AssistantMessage,
     ChatCompletionClient,
@@ -74,7 +73,7 @@ class AssistantAgent(BaseChatAgent):
         .. code-block:: python
 
             import asyncio
-            from autogen_core.base import CancellationToken
+            from autogen_core import CancellationToken
             from autogen_ext.models import OpenAIChatCompletionClient
             from autogen_agentchat.agents import AssistantAgent
             from autogen_agentchat.messages import TextMessage
@@ -107,7 +106,7 @@ class AssistantAgent(BaseChatAgent):
             from autogen_agentchat.agents import AssistantAgent
             from autogen_agentchat.messages import TextMessage
             from autogen_agentchat.task import Console
-            from autogen_core.base import CancellationToken
+            from autogen_core import CancellationToken
 
 
             async def get_current_time() -> str:
@@ -136,7 +135,7 @@ class AssistantAgent(BaseChatAgent):
         .. code-block:: python
 
             import asyncio
-            from autogen_core.base import CancellationToken
+            from autogen_core import CancellationToken
             from autogen_ext.models import OpenAIChatCompletionClient
             from autogen_agentchat.agents import AssistantAgent
             from autogen_agentchat.messages import TextMessage

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_base_chat_agent.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import AsyncGenerator, List, Sequence
 
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken
 
 from ..base import ChatAgent, Response, TaskResult
 from ..messages import AgentMessage, ChatMessage, HandoffMessage, MultiModalMessage, StopMessage, TextMessage

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_code_executor_agent.py

@@ -1,6 +1,6 @@
 from typing import List, Sequence
 
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken
 from autogen_core.components.code_executor import CodeBlock, CodeExecutor, extract_markdown_code_blocks
 
 from ..base import Response
@@ -28,7 +28,7 @@ class CodeExecutorAgent(BaseChatAgent):
             from autogen_agentchat.agents import CodeExecutorAgent
             from autogen_agentchat.messages import TextMessage
             from autogen_ext.code_executors import DockerCommandLineCodeExecutor
-            from autogen_core.base import CancellationToken
+            from autogen_core import CancellationToken
 
 
             async def run_code_executor_agent() -> None:

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py

@@ -1,7 +1,6 @@
 from typing import AsyncGenerator, List, Sequence
 
-from autogen_core.base import CancellationToken
-from autogen_core.components import Image
+from autogen_core import CancellationToken, Image
 from autogen_core.components.models import ChatCompletionClient
 from autogen_core.components.models._types import SystemMessage
 

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py

@@ -2,7 +2,7 @@ import asyncio
 from inspect import iscoroutinefunction
 from typing import Awaitable, Callable, List, Optional, Sequence, Union, cast
 
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken
 
 from ..base import Response
 from ..messages import ChatMessage, HandoffMessage, TextMessage

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

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import AsyncGenerator, List, Protocol, Sequence, runtime_checkable
 
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken
 
 from ..messages import AgentMessage, ChatMessage
 from ._task import TaskRunner

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

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import AsyncGenerator, Protocol, Sequence
 
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken
 
 from ..messages import AgentMessage, ChatMessage
 

python/packages/autogen-agentchat/src/autogen_agentchat/messages.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from autogen_core.components import FunctionCall, Image
+from autogen_core import FunctionCall, Image
 from autogen_core.components.models import FunctionExecutionResult, RequestUsage
 from pydantic import BaseModel, ConfigDict
 

python/packages/autogen-agentchat/src/autogen_agentchat/task/_console.py

@@ -3,7 +3,7 @@ import sys
 import time
 from typing import AsyncGenerator, List, Optional, TypeVar, cast
 
-from autogen_core.components import Image
+from autogen_core import Image
 from autogen_core.components.models import RequestUsage
 
 from autogen_agentchat.base import Response, TaskResult

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

@@ -4,17 +4,18 @@ import uuid
 from abc import ABC, abstractmethod
 from typing import AsyncGenerator, Callable, List
 
-from autogen_core.application import SingleThreadedAgentRuntime
-from autogen_core.base import (
+from autogen_core import (
     AgentId,
     AgentInstantiationContext,
     AgentRuntime,
     AgentType,
     CancellationToken,
+    ClosureAgent,
     MessageContext,
+    TypeSubscription,
 )
-from autogen_core.components import ClosureAgent, TypeSubscription
-from autogen_core.components._closure_agent import ClosureContext
+from autogen_core._closure_agent import ClosureContext
+from autogen_core.application import SingleThreadedAgentRuntime
 
 from ... import EVENT_LOGGER_NAME
 from ...base import ChatAgent, TaskResult, Team, TerminationCondition
@@ -216,7 +217,7 @@ class BaseGroupChat(Team, ABC):
             from autogen_agentchat.agents import AssistantAgent
             from autogen_agentchat.task import MaxMessageTermination
             from autogen_agentchat.teams import RoundRobinGroupChat
-            from autogen_core.base import CancellationToken
+            from autogen_core import CancellationToken
             from autogen_ext.models import OpenAIChatCompletionClient
 
 
@@ -316,7 +317,7 @@ class BaseGroupChat(Team, ABC):
             from autogen_agentchat.agents import AssistantAgent
             from autogen_agentchat.task import MaxMessageTermination, Console
             from autogen_agentchat.teams import RoundRobinGroupChat
-            from autogen_core.base import CancellationToken
+            from autogen_core import CancellationToken
             from autogen_ext.models import OpenAIChatCompletionClient
 
 

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

@@ -2,8 +2,7 @@ import asyncio
 from abc import ABC, abstractmethod
 from typing import Any, List
 
-from autogen_core.base import MessageContext
-from autogen_core.components import DefaultTopicId, event, rpc
+from autogen_core import DefaultTopicId, MessageContext, event, rpc
 
 from ...base import TerminationCondition
 from ...messages import AgentMessage, ChatMessage, StopMessage

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

@@ -1,7 +1,6 @@
 from typing import Any, List
 
-from autogen_core.base import MessageContext
-from autogen_core.components import DefaultTopicId, event, rpc
+from autogen_core import DefaultTopicId, MessageContext, event, rpc
 
 from ...base import ChatAgent, Response
 from ...messages import ChatMessage

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

@@ -2,8 +2,7 @@ import json
 import logging
 from typing import Any, Dict, List
 
-from autogen_core.base import AgentId, CancellationToken, MessageContext
-from autogen_core.components import DefaultTopicId, Image, event, rpc
+from autogen_core import AgentId, CancellationToken, DefaultTopicId, Image, MessageContext, event, rpc
 from autogen_core.components.models import (
     AssistantMessage,
     ChatCompletionClient,

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

@@ -1,8 +1,7 @@
 import asyncio
 from typing import Any
 
-from autogen_core.base import MessageContext
-from autogen_core.components import RoutedAgent
+from autogen_core import MessageContext, RoutedAgent
 
 
 class FIFOLock:

python/packages/autogen-agentchat/tests/test_assistant_agent.py

@@ -14,7 +14,7 @@ from autogen_agentchat.messages import (
     ToolCallMessage,
     ToolCallResultMessage,
 )
-from autogen_core.components import Image
+from autogen_core import Image
 from autogen_core.components.tools import FunctionTool
 from autogen_ext.models import OpenAIChatCompletionClient
 from openai.resources.chat.completions import AsyncCompletions

python/packages/autogen-agentchat/tests/test_group_chat.py

@@ -28,8 +28,7 @@ from autogen_agentchat.teams import (
     SelectorGroupChat,
     Swarm,
 )
-from autogen_core.base import CancellationToken
-from autogen_core.components import FunctionCall
+from autogen_core import CancellationToken, FunctionCall
 from autogen_core.components.code_executor import LocalCommandLineCodeExecutor
 from autogen_core.components.models import FunctionExecutionResult
 from autogen_core.components.tools import FunctionTool

python/packages/autogen-agentchat/tests/test_magentic_one_group_chat.py

@@ -16,7 +16,7 @@ from autogen_agentchat.messages import (
 from autogen_agentchat.teams import (
     MagenticOneGroupChat,
 )
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken
 from autogen_ext.models import ReplayChatCompletionClient
 from utils import FileLogHandler
 

python/packages/autogen-agentchat/tests/test_sequential_routed_agent.py

@@ -5,9 +5,8 @@ from typing import List
 
 import pytest
 from autogen_agentchat.teams._group_chat._sequential_routed_agent import SequentialRoutedAgent
+from autogen_core import AgentId, DefaultTopicId, MessageContext, default_subscription, message_handler
 from autogen_core.application import SingleThreadedAgentRuntime
-from autogen_core.base import AgentId, MessageContext
-from autogen_core.components import DefaultTopicId, default_subscription, message_handler
 
 
 @dataclass

python/packages/autogen-agentchat/tests/test_userproxy_agent.py

@@ -5,7 +5,7 @@ import pytest
 from autogen_agentchat.agents import UserProxyAgent
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import ChatMessage, HandoffMessage, TextMessage
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken
 
 
 @pytest.mark.asyncio

python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/agents.ipynb

@@ -35,7 +35,7 @@
    "source": [
     "from autogen_agentchat.agents import AssistantAgent\n",
     "from autogen_agentchat.messages import TextMessage\n",
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "from autogen_ext.models import OpenAIChatCompletionClient\n",
     "\n",
     "\n",

python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/tutorial/custom-agents.ipynb

@@ -43,7 +43,7 @@
     "from autogen_agentchat.agents import BaseChatAgent\n",
     "from autogen_agentchat.base import Response\n",
     "from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage\n",
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "\n",
     "\n",
     "class CountDownAgent(BaseChatAgent):\n",
@@ -120,7 +120,7 @@
     "from autogen_agentchat.agents import BaseChatAgent\n",
     "from autogen_agentchat.base import Response\n",
     "from autogen_agentchat.messages import ChatMessage\n",
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "\n",
     "\n",
     "class UserProxyAgent(BaseChatAgent):\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/extracting-results-with-an-agent.ipynb

@@ -25,9 +25,8 @@
     "import asyncio\n",
     "from dataclasses import dataclass\n",
     "\n",
-    "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import MessageContext\n",
-    "from autogen_core.components import ClosureAgent, ClosureContext, DefaultSubscription, DefaultTopicId"
+    "from autogen_core import ClosureAgent, ClosureContext, DefaultSubscription, DefaultTopicId, MessageContext\n",
+    "from autogen_core.application import SingleThreadedAgentRuntime"
    ]
   },
   {

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/langgraph-agent.ipynb

@@ -42,9 +42,8 @@
     "from dataclasses import dataclass\n",
     "from typing import Any, Callable, List, Literal\n",
     "\n",
+    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import AgentId, MessageContext\n",
-    "from autogen_core.components import RoutedAgent, message_handler\n",
     "from azure.identity import DefaultAzureCredential, get_bearer_token_provider\n",
     "from langchain_core.messages import HumanMessage, SystemMessage\n",
     "from langchain_core.tools import tool  # pyright: ignore\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/llamaindex-agent.ipynb

@@ -41,9 +41,8 @@
     "from dataclasses import dataclass\n",
     "from typing import List, Optional\n",
     "\n",
+    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import AgentId, MessageContext\n",
-    "from autogen_core.components import RoutedAgent, message_handler\n",
     "from azure.identity import DefaultAzureCredential, get_bearer_token_provider\n",
     "from llama_index.core import Settings\n",
     "from llama_index.core.agent import ReActAgent\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/local-llms-ollama-litellm.ipynb

@@ -39,9 +39,8 @@
    "source": [
     "from dataclasses import dataclass\n",
     "\n",
+    "from autogen_core import AgentId, DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import AgentId, MessageContext\n",
-    "from autogen_core.components import DefaultTopicId, RoutedAgent, default_subscription, message_handler\n",
     "from autogen_core.components.model_context import BufferedChatCompletionContext\n",
     "from autogen_core.components.models import (\n",
     "    AssistantMessage,\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/openai-assistant-agent.ipynb

@@ -104,8 +104,7 @@
     "from typing import Any, Callable, List\n",
     "\n",
     "import aiofiles\n",
-    "from autogen_core.base import AgentId, MessageContext\n",
-    "from autogen_core.components import RoutedAgent, message_handler\n",
+    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
     "from openai import AsyncAssistantEventHandler, AsyncClient\n",
     "from openai.types.beta.thread import ToolResources, ToolResourcesFileSearch\n",
     "\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/termination-with-intervention.ipynb

@@ -22,10 +22,9 @@
     "from dataclasses import dataclass\n",
     "from typing import Any\n",
     "\n",
+    "from autogen_core import AgentId, DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import AgentId, MessageContext\n",
-    "from autogen_core.base.intervention import DefaultInterventionHandler\n",
-    "from autogen_core.components import DefaultTopicId, RoutedAgent, default_subscription, message_handler"
+    "from autogen_core.base.intervention import DefaultInterventionHandler"
    ]
   },
   {

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/tool-use-with-intervention.ipynb

@@ -19,18 +19,17 @@
     "from dataclasses import dataclass\n",
     "from typing import Any, List\n",
     "\n",
+    "from autogen_core import AgentId, AgentType, FunctionCall, MessageContext, RoutedAgent, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import AgentId, AgentType, MessageContext\n",
     "from autogen_core.base.intervention import DefaultInterventionHandler, DropMessage\n",
-    "from autogen_core.components import FunctionCall, RoutedAgent, message_handler\n",
     "from autogen_core.components.models import (\n",
     "    ChatCompletionClient,\n",
     "    LLMMessage,\n",
     "    SystemMessage,\n",
     "    UserMessage,\n",
     ")\n",
-    "from autogen_core.components.tool_agent import ToolAgent, ToolException, tool_agent_caller_loop\n",
     "from autogen_core.components.tools import PythonCodeExecutionTool, ToolSchema\n",
+    "from autogen_core.tool_agent import ToolAgent, ToolException, tool_agent_caller_loop\n",
     "from autogen_ext.code_executors import DockerCommandLineCodeExecutor\n",
     "from autogen_ext.models import OpenAIChatCompletionClient"
    ]

python/packages/autogen-core/docs/src/user-guide/core-user-guide/cookbook/topic-subscription-scenarios.ipynb

@@ -44,11 +44,10 @@
     "from enum import Enum\n",
     "from typing import List\n",
     "\n",
+    "from autogen_core import MessageContext, RoutedAgent, TopicId, TypeSubscription, message_handler\n",
+    "from autogen_core._default_subscription import DefaultSubscription\n",
+    "from autogen_core._default_topic import DefaultTopicId\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import MessageContext, TopicId\n",
-    "from autogen_core.components import RoutedAgent, TypeSubscription, message_handler\n",
-    "from autogen_core.components._default_subscription import DefaultSubscription\n",
-    "from autogen_core.components._default_topic import DefaultTopicId\n",
     "from autogen_core.components.models import (\n",
     "    SystemMessage,\n",
     ")"

python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/group-chat.ipynb

@@ -72,16 +72,17 @@
     "from typing import List\n",
     "\n",
     "import openai\n",
-    "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import MessageContext, TopicId\n",
-    "from autogen_core.components import (\n",
+    "from autogen_core import (\n",
     "    DefaultTopicId,\n",
     "    FunctionCall,\n",
     "    Image,\n",
+    "    MessageContext,\n",
     "    RoutedAgent,\n",
+    "    TopicId,\n",
     "    TypeSubscription,\n",
     "    message_handler,\n",
     ")\n",
+    "from autogen_core.application import SingleThreadedAgentRuntime\n",
     "from autogen_core.components.models import (\n",
     "    AssistantMessage,\n",
     "    ChatCompletionClient,\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/handoffs.ipynb

@@ -56,9 +56,8 @@
     "import uuid\n",
     "from typing import List, Tuple\n",
     "\n",
+    "from autogen_core import FunctionCall, MessageContext, RoutedAgent, TopicId, TypeSubscription, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import MessageContext, TopicId\n",
-    "from autogen_core.components import FunctionCall, RoutedAgent, TypeSubscription, message_handler\n",
     "from autogen_core.components.models import (\n",
     "    AssistantMessage,\n",
     "    ChatCompletionClient,\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/mixture-of-agents.ipynb

@@ -38,9 +38,8 @@
     "from dataclasses import dataclass\n",
     "from typing import List\n",
     "\n",
+    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import AgentId, MessageContext\n",
-    "from autogen_core.components import RoutedAgent, message_handler\n",
     "from autogen_core.components.models import ChatCompletionClient, SystemMessage, UserMessage\n",
     "from autogen_ext.models import OpenAIChatCompletionClient"
    ]

python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/multi-agent-debate.ipynb

@@ -43,9 +43,15 @@
     "from dataclasses import dataclass\n",
     "from typing import Dict, List\n",
     "\n",
+    "from autogen_core import (\n",
+    "    DefaultTopicId,\n",
+    "    MessageContext,\n",
+    "    RoutedAgent,\n",
+    "    TypeSubscription,\n",
+    "    default_subscription,\n",
+    "    message_handler,\n",
+    ")\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import MessageContext\n",
-    "from autogen_core.components import DefaultTopicId, RoutedAgent, TypeSubscription, default_subscription, message_handler\n",
     "from autogen_core.components.models import (\n",
     "    AssistantMessage,\n",
     "    ChatCompletionClient,\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/design-patterns/reflection.ipynb

@@ -100,8 +100,7 @@
     "import uuid\n",
     "from typing import Dict, List, Union\n",
     "\n",
-    "from autogen_core.base import MessageContext, TopicId\n",
-    "from autogen_core.components import RoutedAgent, default_subscription, message_handler\n",
+    "from autogen_core import MessageContext, RoutedAgent, TopicId, default_subscription, message_handler\n",
     "from autogen_core.components.models import (\n",
     "    AssistantMessage,\n",
     "    ChatCompletionClient,\n",
@@ -442,8 +441,8 @@
     }
    ],
    "source": [
+    "from autogen_core import DefaultTopicId\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.components import DefaultTopicId\n",
     "from autogen_ext.models import OpenAIChatCompletionClient\n",
     "\n",
     "runtime = SingleThreadedAgentRuntime()\n",

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

@@ -55,7 +55,7 @@
    "source": [
     "from dataclasses import dataclass\n",
     "\n",
-    "from autogen_core.base import AgentId, BaseAgent, MessageContext\n",
+    "from autogen_core import AgentId, BaseAgent, MessageContext\n",
     "\n",
     "\n",
     "@dataclass\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/command-line-code-executors.ipynb

@@ -54,7 +54,7 @@
    "source": [
     "from pathlib import Path\n",
     "\n",
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "from autogen_core.components.code_executor import CodeBlock\n",
     "from autogen_ext.code_executors import DockerCommandLineCodeExecutor\n",
     "\n",
@@ -125,7 +125,7 @@
    "source": [
     "from pathlib import Path\n",
     "\n",
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "from autogen_core.components.code_executor import CodeBlock, LocalCommandLineCodeExecutor\n",
     "\n",
     "work_dir = Path(\"coding\")\n",
@@ -171,7 +171,7 @@
     "import venv\n",
     "from pathlib import Path\n",
     "\n",
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "from autogen_core.components.code_executor import CodeBlock, LocalCommandLineCodeExecutor\n",
     "\n",
     "work_dir = Path(\"coding\")\n",

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

@@ -1,223 +1,222 @@
 {
-    "cells": [
-        {
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "# Distributed Agent Runtime\n",
-                "\n",
-                "```{attention}\n",
-                "The distributed agent runtime is an experimental feature. Expect breaking changes\n",
-                "to the API.\n",
-                "```\n",
-                "\n",
-                "A distributed agent runtime facilitates communication and agent lifecycle management\n",
-                "across process boundaries.\n",
-                "It consists of a host service and at least one worker runtime.\n",
-                "\n",
-                "The host service maintains connections to all active worker runtimes,\n",
-                "facilitates message delivery, and keeps sessions for all direct messages (i.e., RPCs).\n",
-                "A worker runtime processes application code (agents) and connects to the host service.\n",
-                "It also advertises the agents which they support to the host service,\n",
-                "so the host service can deliver messages to the correct worker.\n",
-                "\n",
-                "````{note}\n",
-                "The distributed agent runtime requires extra dependencies, install them using:\n",
-                "```bash\n",
-                "pip install autogen-core[grpc]==0.4.0.dev8\n",
-                "```\n",
-                "````\n",
-                "\n",
-                "We can start a host service using {py:class}`~autogen_core.application.WorkerAgentRuntimeHost`."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 1,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "from autogen_core.application import WorkerAgentRuntimeHost\n",
-                "\n",
-                "host = WorkerAgentRuntimeHost(address=\"localhost:50051\")\n",
-                "host.start()  # Start a host service in the background."
-            ]
-        },
-        {
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "The above code starts the host service in the background and accepts\n",
-                "worker connections on port 50051.\n",
-                "\n",
-                "Before running worker runtimes, let's define our agent.\n",
-                "The agent will publish a new message on every message it receives.\n",
-                "It also keeps track of how many messages it has published, and \n",
-                "stops publishing new messages once it has published 5 messages."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 2,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "from dataclasses import dataclass\n",
-                "\n",
-                "from autogen_core.base import MessageContext\n",
-                "from autogen_core.components import DefaultTopicId, RoutedAgent, default_subscription, message_handler\n",
-                "\n",
-                "\n",
-                "@dataclass\n",
-                "class MyMessage:\n",
-                "    content: str\n",
-                "\n",
-                "\n",
-                "@default_subscription\n",
-                "class MyAgent(RoutedAgent):\n",
-                "    def __init__(self, name: str) -> None:\n",
-                "        super().__init__(\"My agent\")\n",
-                "        self._name = name\n",
-                "        self._counter = 0\n",
-                "\n",
-                "    @message_handler\n",
-                "    async def my_message_handler(self, message: MyMessage, ctx: MessageContext) -> None:\n",
-                "        self._counter += 1\n",
-                "        if self._counter > 5:\n",
-                "            return\n",
-                "        content = f\"{self._name}: Hello x {self._counter}\"\n",
-                "        print(content)\n",
-                "        await self.publish_message(MyMessage(content=content), DefaultTopicId())"
-            ]
-        },
-        {
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "Now we can set up the worker agent runtimes.\n",
-                "We use {py:class}`~autogen_core.application.WorkerAgentRuntime`.\n",
-                "We set up two worker runtimes. Each runtime hosts one agent.\n",
-                "All agents publish and subscribe to the default topic, so they can see all\n",
-                "messages being published.\n",
-                "\n",
-                "To run the agents, we publishes a message from a worker."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 3,
-            "metadata": {},
-            "outputs": [
-                {
-                    "name": "stdout",
-                    "output_type": "stream",
-                    "text": [
-                        "worker1: Hello x 1\n",
-                        "worker2: Hello x 1\n",
-                        "worker2: Hello x 2\n",
-                        "worker1: Hello x 2\n",
-                        "worker1: Hello x 3\n",
-                        "worker2: Hello x 3\n",
-                        "worker2: Hello x 4\n",
-                        "worker1: Hello x 4\n",
-                        "worker1: Hello x 5\n",
-                        "worker2: Hello x 5\n"
-                    ]
-                }
-            ],
-            "source": [
-                "import asyncio\n",
-                "\n",
-                "from autogen_core.application import WorkerAgentRuntime\n",
-                "\n",
-                "worker1 = WorkerAgentRuntime(host_address=\"localhost:50051\")\n",
-                "worker1.start()\n",
-                "await MyAgent.register(worker1, \"worker1\", lambda: MyAgent(\"worker1\"))\n",
-                "\n",
-                "worker2 = WorkerAgentRuntime(host_address=\"localhost:50051\")\n",
-                "worker2.start()\n",
-                "await MyAgent.register(worker2, \"worker2\", lambda: MyAgent(\"worker2\"))\n",
-                "\n",
-                "await worker2.publish_message(MyMessage(content=\"Hello!\"), DefaultTopicId())\n",
-                "\n",
-                "# Let the agents run for a while.\n",
-                "await asyncio.sleep(5)"
-            ]
-        },
-        {
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "We can see each agent published exactly 5 messages.\n",
-                "\n",
-                "To stop the worker runtimes, we can call {py:meth}`~autogen_core.application.WorkerAgentRuntime.stop`."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 4,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "await worker1.stop()\n",
-                "await worker2.stop()\n",
-                "\n",
-                "# To keep the worker running until a termination signal is received (e.g., SIGTERM).\n",
-                "# await worker1.stop_when_signal()"
-            ]
-        },
-        {
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "We can call {py:meth}`~autogen_core.application.WorkerAgentRuntimeHost.stop`\n",
-                "to stop the host service."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 5,
-            "metadata": {},
-            "outputs": [],
-            "source": [
-                "await host.stop()\n",
-                "\n",
-                "# To keep the host service running until a termination signal (e.g., SIGTERM)\n",
-                "# await host.stop_when_signal()"
-            ]
-        },
-        {
-            "cell_type": "markdown",
-            "metadata": {},
-            "source": [
-                "# Next Steps\n",
-                "To see complete examples of using distributed runtime, please take a look at the following samples:\n",
-                "\n",
-                "- [Distributed Workers](https://github.com/microsoft/autogen/tree/main/python/packages/autogen-core/samples/worker)  \n",
-                "- [Distributed Semantic Router](https://github.com/microsoft/autogen/tree/main/python/packages/autogen-core/samples/semantic_router)  \n",
-                "- [Distributed Group Chat](https://github.com/microsoft/autogen/tree/main/python/packages/autogen-core/samples/distributed-group-chat)  \n"
-            ]
-        }
-    ],
-    "metadata": {
-        "kernelspec": {
-            "display_name": "agnext",
-            "language": "python",
-            "name": "python3"
-        },
-        "language_info": {
-            "codemirror_mode": {
-                "name": "ipython",
-                "version": 3
-            },
-            "file_extension": ".py",
-            "mimetype": "text/x-python",
-            "name": "python",
-            "nbconvert_exporter": "python",
-            "pygments_lexer": "ipython3",
-            "version": "3.11.9"
-        }
-    },
-    "nbformat": 4,
-    "nbformat_minor": 2
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Distributed Agent Runtime\n",
+    "\n",
+    "```{attention}\n",
+    "The distributed agent runtime is an experimental feature. Expect breaking changes\n",
+    "to the API.\n",
+    "```\n",
+    "\n",
+    "A distributed agent runtime facilitates communication and agent lifecycle management\n",
+    "across process boundaries.\n",
+    "It consists of a host service and at least one worker runtime.\n",
+    "\n",
+    "The host service maintains connections to all active worker runtimes,\n",
+    "facilitates message delivery, and keeps sessions for all direct messages (i.e., RPCs).\n",
+    "A worker runtime processes application code (agents) and connects to the host service.\n",
+    "It also advertises the agents which they support to the host service,\n",
+    "so the host service can deliver messages to the correct worker.\n",
+    "\n",
+    "````{note}\n",
+    "The distributed agent runtime requires extra dependencies, install them using:\n",
+    "```bash\n",
+    "pip install autogen-core[grpc]==0.4.0.dev8\n",
+    "```\n",
+    "````\n",
+    "\n",
+    "We can start a host service using {py:class}`~autogen_core.application.WorkerAgentRuntimeHost`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from autogen_core.application import WorkerAgentRuntimeHost\n",
+    "\n",
+    "host = WorkerAgentRuntimeHost(address=\"localhost:50051\")\n",
+    "host.start()  # Start a host service in the background."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The above code starts the host service in the background and accepts\n",
+    "worker connections on port 50051.\n",
+    "\n",
+    "Before running worker runtimes, let's define our agent.\n",
+    "The agent will publish a new message on every message it receives.\n",
+    "It also keeps track of how many messages it has published, and \n",
+    "stops publishing new messages once it has published 5 messages."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dataclasses import dataclass\n",
+    "\n",
+    "from autogen_core import DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler\n",
+    "\n",
+    "\n",
+    "@dataclass\n",
+    "class MyMessage:\n",
+    "    content: str\n",
+    "\n",
+    "\n",
+    "@default_subscription\n",
+    "class MyAgent(RoutedAgent):\n",
+    "    def __init__(self, name: str) -> None:\n",
+    "        super().__init__(\"My agent\")\n",
+    "        self._name = name\n",
+    "        self._counter = 0\n",
+    "\n",
+    "    @message_handler\n",
+    "    async def my_message_handler(self, message: MyMessage, ctx: MessageContext) -> None:\n",
+    "        self._counter += 1\n",
+    "        if self._counter > 5:\n",
+    "            return\n",
+    "        content = f\"{self._name}: Hello x {self._counter}\"\n",
+    "        print(content)\n",
+    "        await self.publish_message(MyMessage(content=content), DefaultTopicId())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can set up the worker agent runtimes.\n",
+    "We use {py:class}`~autogen_core.application.WorkerAgentRuntime`.\n",
+    "We set up two worker runtimes. Each runtime hosts one agent.\n",
+    "All agents publish and subscribe to the default topic, so they can see all\n",
+    "messages being published.\n",
+    "\n",
+    "To run the agents, we publishes a message from a worker."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "worker1: Hello x 1\n",
+      "worker2: Hello x 1\n",
+      "worker2: Hello x 2\n",
+      "worker1: Hello x 2\n",
+      "worker1: Hello x 3\n",
+      "worker2: Hello x 3\n",
+      "worker2: Hello x 4\n",
+      "worker1: Hello x 4\n",
+      "worker1: Hello x 5\n",
+      "worker2: Hello x 5\n"
+     ]
+    }
+   ],
+   "source": [
+    "import asyncio\n",
+    "\n",
+    "from autogen_core.application import WorkerAgentRuntime\n",
+    "\n",
+    "worker1 = WorkerAgentRuntime(host_address=\"localhost:50051\")\n",
+    "worker1.start()\n",
+    "await MyAgent.register(worker1, \"worker1\", lambda: MyAgent(\"worker1\"))\n",
+    "\n",
+    "worker2 = WorkerAgentRuntime(host_address=\"localhost:50051\")\n",
+    "worker2.start()\n",
+    "await MyAgent.register(worker2, \"worker2\", lambda: MyAgent(\"worker2\"))\n",
+    "\n",
+    "await worker2.publish_message(MyMessage(content=\"Hello!\"), DefaultTopicId())\n",
+    "\n",
+    "# Let the agents run for a while.\n",
+    "await asyncio.sleep(5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can see each agent published exactly 5 messages.\n",
+    "\n",
+    "To stop the worker runtimes, we can call {py:meth}`~autogen_core.application.WorkerAgentRuntime.stop`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "await worker1.stop()\n",
+    "await worker2.stop()\n",
+    "\n",
+    "# To keep the worker running until a termination signal is received (e.g., SIGTERM).\n",
+    "# await worker1.stop_when_signal()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can call {py:meth}`~autogen_core.application.WorkerAgentRuntimeHost.stop`\n",
+    "to stop the host service."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "await host.stop()\n",
+    "\n",
+    "# To keep the host service running until a termination signal (e.g., SIGTERM)\n",
+    "# await host.stop_when_signal()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Next Steps\n",
+    "To see complete examples of using distributed runtime, please take a look at the following samples:\n",
+    "\n",
+    "- [Distributed Workers](https://github.com/microsoft/autogen/tree/main/python/packages/autogen-core/samples/worker)  \n",
+    "- [Distributed Semantic Router](https://github.com/microsoft/autogen/tree/main/python/packages/autogen-core/samples/semantic_router)  \n",
+    "- [Distributed Group Chat](https://github.com/microsoft/autogen/tree/main/python/packages/autogen-core/samples/distributed-group-chat)  \n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "agnext",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
 }

python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/message-and-communication.ipynb

@@ -90,9 +90,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import AgentId, MessageContext\n",
-    "from autogen_core.components import RoutedAgent, message_handler\n",
     "\n",
     "\n",
     "class MyAgent(RoutedAgent):\n",
@@ -299,9 +298,8 @@
    "source": [
     "from dataclasses import dataclass\n",
     "\n",
+    "from autogen_core import MessageContext, RoutedAgent, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import MessageContext\n",
-    "from autogen_core.components import RoutedAgent, message_handler\n",
     "\n",
     "\n",
     "@dataclass\n",
@@ -420,7 +418,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from autogen_core.components import RoutedAgent, message_handler, type_subscription\n",
+    "from autogen_core import RoutedAgent, message_handler, type_subscription\n",
     "\n",
     "\n",
     "@type_subscription(topic_type=\"default\")\n",
@@ -452,7 +450,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from autogen_core.base import TopicId\n",
+    "from autogen_core import TopicId\n",
     "\n",
     "\n",
     "class BroadcastingAgent(RoutedAgent):\n",
@@ -498,7 +496,7 @@
     }
    ],
    "source": [
-    "from autogen_core.components import TypeSubscription\n",
+    "from autogen_core import TypeSubscription\n",
     "\n",
     "runtime = SingleThreadedAgentRuntime()\n",
     "\n",
@@ -561,7 +559,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from autogen_core.components import DefaultTopicId, default_subscription\n",
+    "from autogen_core import DefaultTopicId, default_subscription\n",
     "\n",
     "\n",
     "@default_subscription\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/model-clients.ipynb

@@ -329,9 +329,8 @@
    "source": [
     "from dataclasses import dataclass\n",
     "\n",
+    "from autogen_core import MessageContext, RoutedAgent, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import MessageContext\n",
-    "from autogen_core.components import RoutedAgent, message_handler\n",
     "from autogen_core.components.models import ChatCompletionClient, SystemMessage, UserMessage\n",
     "from autogen_ext.models import OpenAIChatCompletionClient\n",
     "\n",
@@ -422,7 +421,7 @@
    ],
    "source": [
     "# Create the runtime and register the agent.\n",
-    "from autogen_core.base import AgentId\n",
+    "from autogen_core import AgentId\n",
     "\n",
     "runtime = SingleThreadedAgentRuntime()\n",
     "await SimpleAgent.register(\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/framework/tools.ipynb

@@ -43,7 +43,7 @@
     }
    ],
    "source": [
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "from autogen_core.components.tools import PythonCodeExecutionTool\n",
     "from autogen_ext.code_executors import DockerCommandLineCodeExecutor\n",
     "\n",
@@ -113,7 +113,7 @@
    "source": [
     "import random\n",
     "\n",
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "from autogen_core.components.tools import FunctionTool\n",
     "from typing_extensions import Annotated\n",
     "\n",
@@ -155,17 +155,16 @@
     "from dataclasses import dataclass\n",
     "from typing import List\n",
     "\n",
+    "from autogen_core import AgentId, AgentInstantiationContext, MessageContext, RoutedAgent, message_handler\n",
     "from autogen_core.application import SingleThreadedAgentRuntime\n",
-    "from autogen_core.base import AgentId, AgentInstantiationContext, MessageContext\n",
-    "from autogen_core.components import RoutedAgent, message_handler\n",
     "from autogen_core.components.models import (\n",
     "    ChatCompletionClient,\n",
     "    LLMMessage,\n",
     "    SystemMessage,\n",
     "    UserMessage,\n",
     ")\n",
-    "from autogen_core.components.tool_agent import ToolAgent, tool_agent_caller_loop\n",
     "from autogen_core.components.tools import FunctionTool, Tool, ToolSchema\n",
+    "from autogen_core.tool_agent import ToolAgent, tool_agent_caller_loop\n",
     "from autogen_ext.models import OpenAIChatCompletionClient\n",
     "\n",
     "\n",

python/packages/autogen-core/docs/src/user-guide/core-user-guide/quickstart.ipynb

@@ -34,8 +34,7 @@
     "from dataclasses import dataclass\n",
     "from typing import List\n",
     "\n",
-    "from autogen_core.base import MessageContext\n",
-    "from autogen_core.components import DefaultTopicId, RoutedAgent, default_subscription, message_handler\n",
+    "from autogen_core import DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler\n",
     "from autogen_core.components.code_executor import CodeExecutor, extract_markdown_code_blocks\n",
     "from autogen_core.components.models import (\n",
     "    AssistantMessage,\n",

python/packages/autogen-core/docs/src/user-guide/extensions-user-guide/azure-container-code-executor.ipynb

@@ -61,7 +61,7 @@
     "import tempfile\n",
     "\n",
     "from anyio import open_file\n",
-    "from autogen_core.base import CancellationToken\n",
+    "from autogen_core import CancellationToken\n",
     "from autogen_core.components.code_executor import CodeBlock\n",
     "from autogen_ext.code_executor.aca_dynamic_sessions import AzureContainerCodeExecutor\n",
     "from azure.identity import DefaultAzureCredential"

python/packages/autogen-core/samples/chess_game.py

@@ -7,9 +7,8 @@ import asyncio
 import logging
 from typing import Annotated, Literal
 
+from autogen_core import AgentId, AgentInstantiationContext, AgentRuntime, DefaultSubscription, DefaultTopicId
 from autogen_core.application import SingleThreadedAgentRuntime
-from autogen_core.base import AgentId, AgentInstantiationContext, AgentRuntime
-from autogen_core.components import DefaultSubscription, DefaultTopicId
 from autogen_core.components.model_context import BufferedChatCompletionContext
 from autogen_core.components.models import SystemMessage
 from autogen_core.components.tools import FunctionTool

python/packages/autogen-core/samples/common/agents/_chat_completion_agent.py

@@ -2,10 +2,12 @@ import asyncio
 import json
 from typing import Any, Coroutine, Dict, List, Mapping, Sequence, Tuple
 
-from autogen_core.base import AgentId, CancellationToken, MessageContext
-from autogen_core.components import (
+from autogen_core import (
+    AgentId,
+    CancellationToken,
     DefaultTopicId,
     FunctionCall,
+    MessageContext,
     RoutedAgent,
     message_handler,
 )

python/packages/autogen-core/samples/common/patterns/_group_chat_manager.py

@@ -1,8 +1,7 @@
 import logging
 from typing import Any, Callable, List, Mapping
 
-from autogen_core.base import AgentId, AgentProxy, MessageContext
-from autogen_core.components import RoutedAgent, message_handler
+from autogen_core import AgentId, AgentProxy, MessageContext, RoutedAgent, message_handler
 from autogen_core.components.model_context import ChatCompletionContext
 from autogen_core.components.models import ChatCompletionClient, UserMessage
 

python/packages/autogen-core/samples/common/patterns/_group_chat_utils.py

@@ -3,7 +3,7 @@
 import re
 from typing import Dict, List
 
-from autogen_core.base import AgentProxy
+from autogen_core import AgentProxy
 from autogen_core.components.model_context import ChatCompletionContext
 from autogen_core.components.models import ChatCompletionClient, SystemMessage, UserMessage
 

python/packages/autogen-core/samples/common/types.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass, field
 from enum import Enum
 from typing import List, Union
 
-from autogen_core.components import FunctionCall, Image
+from autogen_core import FunctionCall, Image
 from autogen_core.components.models import FunctionExecutionResultMessage
 
 

python/packages/autogen-core/samples/distributed-group-chat/_agents.py

@@ -4,9 +4,8 @@ from typing import Awaitable, Callable, List
 from uuid import uuid4
 
 from _types import GroupChatMessage, MessageChunk, RequestToSpeak, UIAgentConfig
+from autogen_core import DefaultTopicId, MessageContext, RoutedAgent, message_handler
 from autogen_core.application import WorkerAgentRuntime
-from autogen_core.base import MessageContext
-from autogen_core.components import DefaultTopicId, RoutedAgent, message_handler
 from autogen_core.components.models import (
     AssistantMessage,
     ChatCompletionClient,

python/packages/autogen-core/samples/distributed-group-chat/_utils.py

@@ -4,7 +4,7 @@ from typing import Any, Iterable, Type
 
 import yaml
 from _types import AppConfig
-from autogen_core.base import MessageSerializer, try_get_known_serializers_for_type
+from autogen_core import MessageSerializer, try_get_known_serializers_for_type
 from autogen_ext.models import AzureOpenAIClientConfiguration
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 

python/packages/autogen-core/samples/distributed-group-chat/run_editor_agent.py

@@ -5,10 +5,10 @@ import warnings
 from _agents import BaseGroupChatAgent
 from _types import AppConfig, GroupChatMessage, MessageChunk, RequestToSpeak
 from _utils import get_serializers, load_config, set_all_log_levels
-from autogen_core.application import WorkerAgentRuntime
-from autogen_core.components import (
+from autogen_core import (
     TypeSubscription,
 )
+from autogen_core.application import WorkerAgentRuntime
 from autogen_ext.models import AzureOpenAIChatCompletionClient
 from rich.console import Console
 from rich.markdown import Markdown

python/packages/autogen-core/samples/distributed-group-chat/run_group_chat_manager.py

@@ -5,10 +5,10 @@ import warnings
 from _agents import GroupChatManager, publish_message_to_ui, publish_message_to_ui_and_backend
 from _types import AppConfig, GroupChatMessage, MessageChunk, RequestToSpeak
 from _utils import get_serializers, load_config, set_all_log_levels
-from autogen_core.application import WorkerAgentRuntime
-from autogen_core.components import (
+from autogen_core import (
     TypeSubscription,
 )
+from autogen_core.application import WorkerAgentRuntime
 from autogen_ext.models import AzureOpenAIChatCompletionClient
 from rich.console import Console
 from rich.markdown import Markdown

python/packages/autogen-core/samples/distributed-group-chat/run_ui.py

@@ -6,10 +6,10 @@ import chainlit as cl  # type: ignore [reportUnknownMemberType] # This dependenc
 from _agents import MessageChunk, UIAgent
 from _types import AppConfig, GroupChatMessage, RequestToSpeak
 from _utils import get_serializers, load_config, set_all_log_levels
-from autogen_core.application import WorkerAgentRuntime
-from autogen_core.components import (
+from autogen_core import (
     TypeSubscription,
 )
+from autogen_core.application import WorkerAgentRuntime
 from chainlit import Message  # type: ignore [reportAttributeAccessIssue]
 from rich.console import Console
 from rich.markdown import Markdown

python/packages/autogen-core/samples/distributed-group-chat/run_writer_agent.py

@@ -5,10 +5,10 @@ import warnings
 from _agents import BaseGroupChatAgent
 from _types import AppConfig, GroupChatMessage, MessageChunk, RequestToSpeak
 from _utils import get_serializers, load_config, set_all_log_levels
-from autogen_core.application import WorkerAgentRuntime
-from autogen_core.components import (
+from autogen_core import (
     TypeSubscription,
 )
+from autogen_core.application import WorkerAgentRuntime
 from autogen_ext.models import AzureOpenAIChatCompletionClient
 from rich.console import Console
 from rich.markdown import Markdown

python/packages/autogen-core/samples/semantic_router/_agents.py

@@ -2,9 +2,8 @@ import asyncio
 import logging
 
 from _semantic_router_components import FinalResult, TerminationMessage, UserProxyMessage, WorkerAgentMessage
+from autogen_core import DefaultTopicId, MessageContext, RoutedAgent, message_handler
 from autogen_core.application.logging import TRACE_LOGGER_NAME
-from autogen_core.base import MessageContext
-from autogen_core.components import DefaultTopicId, RoutedAgent, message_handler
 
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(f"{TRACE_LOGGER_NAME}.workers")

python/packages/autogen-core/samples/semantic_router/_semantic_router_agent.py

@@ -1,9 +1,8 @@
 import logging
 
 from _semantic_router_components import AgentRegistryBase, IntentClassifierBase, TerminationMessage, UserProxyMessage
+from autogen_core import DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler
 from autogen_core.application.logging import TRACE_LOGGER_NAME
-from autogen_core.base import MessageContext
-from autogen_core.components import DefaultTopicId, RoutedAgent, default_subscription, message_handler
 
 logging.basicConfig(level=logging.WARNING)
 logger = logging.getLogger(f"{TRACE_LOGGER_NAME}.semantic_router")

python/packages/autogen-core/samples/semantic_router/run_semantic_router.py

@@ -31,9 +31,8 @@ from _semantic_router_components import (
     UserProxyMessage,
     WorkerAgentMessage,
 )
+from autogen_core import ClosureAgent, ClosureContext, DefaultSubscription, DefaultTopicId, MessageContext
 from autogen_core.application import WorkerAgentRuntime
-from autogen_core.base import MessageContext
-from autogen_core.components import ClosureAgent, ClosureContext, DefaultSubscription, DefaultTopicId
 
 
 class MockIntentClassifier(IntentClassifierBase):

python/packages/autogen-core/samples/slow_human_in_loop.py

@@ -30,16 +30,18 @@ from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from typing import Any, Mapping, Optional
 
-from autogen_core.application import SingleThreadedAgentRuntime
-from autogen_core.base import AgentId, CancellationToken, MessageContext
-from autogen_core.base.intervention import DefaultInterventionHandler
-from autogen_core.components import (
+from autogen_core import (
+    AgentId,
+    CancellationToken,
     DefaultTopicId,
     FunctionCall,
+    MessageContext,
     RoutedAgent,
     message_handler,
     type_subscription,
 )
+from autogen_core.application import SingleThreadedAgentRuntime
+from autogen_core.base.intervention import DefaultInterventionHandler
 from autogen_core.components.model_context import BufferedChatCompletionContext
 from autogen_core.components.models import (
     AssistantMessage,

python/packages/autogen-core/samples/worker/agents.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass
 
-from autogen_core.base import MessageContext
-from autogen_core.components import DefaultTopicId, RoutedAgent, default_subscription, message_handler
+from autogen_core import DefaultTopicId, MessageContext, RoutedAgent, default_subscription, message_handler
 
 
 @dataclass

python/packages/autogen-core/samples/worker/run_cascading_publisher.py

@@ -1,7 +1,6 @@
 from agents import CascadingMessage, ObserverAgent
+from autogen_core import DefaultTopicId, try_get_known_serializers_for_type
 from autogen_core.application import WorkerAgentRuntime
-from autogen_core.base import try_get_known_serializers_for_type
-from autogen_core.components import DefaultTopicId
 
 
 async def main() -> None:

python/packages/autogen-core/samples/worker/run_cascading_worker.py

@@ -1,8 +1,8 @@
 import uuid
 
 from agents import CascadingAgent, ReceiveMessageEvent
+from autogen_core import try_get_known_serializers_for_type
 from autogen_core.application import WorkerAgentRuntime
-from autogen_core.base import try_get_known_serializers_for_type
 
 
 async def main() -> None:

python/packages/autogen-core/samples/worker/run_worker_pub_sub.py

@@ -3,9 +3,15 @@ import logging
 from dataclasses import dataclass
 from typing import Any, NoReturn
 
+from autogen_core import (
+    DefaultSubscription,
+    DefaultTopicId,
+    MessageContext,
+    RoutedAgent,
+    message_handler,
+    try_get_known_serializers_for_type,
+)
 from autogen_core.application import WorkerAgentRuntime
-from autogen_core.base import MessageContext, try_get_known_serializers_for_type
-from autogen_core.components import DefaultSubscription, DefaultTopicId, RoutedAgent, message_handler
 
 
 @dataclass

python/packages/autogen-core/samples/worker/run_worker_rpc.py

@@ -2,12 +2,15 @@ import asyncio
 import logging
 from dataclasses import dataclass
 
-from autogen_core.application import WorkerAgentRuntime
-from autogen_core.base import (
+from autogen_core import (
     AgentId,
+    DefaultSubscription,
+    DefaultTopicId,
     MessageContext,
+    RoutedAgent,
+    message_handler,
 )
-from autogen_core.components import DefaultSubscription, DefaultTopicId, RoutedAgent, message_handler
+from autogen_core.application import WorkerAgentRuntime
 
 
 @dataclass

python/packages/autogen-core/samples/xlang/hello_python_agent/hello_python_agent.py

@@ -3,11 +3,16 @@ import logging
 import os
 import sys
 
-from autogen_core.application import WorkerAgentRuntime
-
 # from protos.agents_events_pb2 import NewMessageReceived
-from autogen_core.base import PROTOBUF_DATA_CONTENT_TYPE, AgentId, try_get_known_serializers_for_type
-from autogen_core.components import DefaultSubscription, DefaultTopicId, TypeSubscription
+from autogen_core import (
+    PROTOBUF_DATA_CONTENT_TYPE,
+    AgentId,
+    DefaultSubscription,
+    DefaultTopicId,
+    TypeSubscription,
+    try_get_known_serializers_for_type,
+)
+from autogen_core.application import WorkerAgentRuntime
 
 # Add the local package directory to sys.path
 thisdir = os.path.dirname(os.path.abspath(__file__))

python/packages/autogen-core/samples/xlang/hello_python_agent/user_input.py

@@ -2,8 +2,7 @@ import asyncio
 import logging
 from typing import Union
 
-from autogen_core.base import MessageContext
-from autogen_core.components import DefaultTopicId, RoutedAgent, message_handler
+from autogen_core import DefaultTopicId, MessageContext, RoutedAgent, message_handler
 from protos.agent_events_pb2 import ConversationClosed, Input, NewMessageReceived, Output  # type: ignore
 
 input_types = Union[ConversationClosed, Input, Output]

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

@@ -1,3 +1,69 @@
 import importlib.metadata
 
 __version__ = importlib.metadata.version("autogen_core")
+
+from ._agent import Agent
+from ._agent_id import AgentId
+from ._agent_instantiation import AgentInstantiationContext
+from ._agent_metadata import AgentMetadata
+from ._agent_proxy import AgentProxy
+from ._agent_runtime import AgentRuntime
+from ._agent_type import AgentType
+from ._base_agent import BaseAgent
+from ._cancellation_token import CancellationToken
+from ._closure_agent import ClosureAgent, ClosureContext
+from ._default_subscription import DefaultSubscription, default_subscription, type_subscription
+from ._default_topic import DefaultTopicId
+from ._image import Image
+from ._message_context import MessageContext
+from ._message_handler_context import MessageHandlerContext
+from ._routed_agent import RoutedAgent, event, message_handler, rpc
+from ._serialization import (
+    JSON_DATA_CONTENT_TYPE,
+    PROTOBUF_DATA_CONTENT_TYPE,
+    MessageSerializer,
+    UnknownPayload,
+    try_get_known_serializers_for_type,
+)
+from ._subscription import Subscription
+from ._subscription_context import SubscriptionInstantiationContext
+from ._topic import TopicId
+from ._type_prefix_subscription import TypePrefixSubscription
+from ._type_subscription import TypeSubscription
+from ._types import FunctionCall
+
+__all__ = [
+    "Agent",
+    "AgentId",
+    "AgentProxy",
+    "AgentMetadata",
+    "AgentRuntime",
+    "BaseAgent",
+    "CancellationToken",
+    "AgentInstantiationContext",
+    "TopicId",
+    "Subscription",
+    "MessageContext",
+    "AgentType",
+    "SubscriptionInstantiationContext",
+    "MessageHandlerContext",
+    "MessageSerializer",
+    "try_get_known_serializers_for_type",
+    "UnknownPayload",
+    "Image",
+    "RoutedAgent",
+    "ClosureAgent",
+    "ClosureContext",
+    "message_handler",
+    "event",
+    "rpc",
+    "FunctionCall",
+    "TypeSubscription",
+    "DefaultSubscription",
+    "DefaultTopicId",
+    "default_subscription",
+    "type_subscription",
+    "TypePrefixSubscription",
+    "JSON_DATA_CONTENT_TYPE",
+    "PROTOBUF_DATA_CONTENT_TYPE",
+]

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

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, Awaitable, Callable, List, Mapping, Protocol, Sequence, TypeVar, get_type_hints
+
+from ._agent_id import AgentId
+from ._agent_instantiation import AgentInstantiationContext
+from ._agent_metadata import AgentMetadata
+from ._agent_runtime import AgentRuntime
+from ._agent_type import AgentType
+from ._base_agent import BaseAgent
+from ._cancellation_token import CancellationToken
+from ._message_context import MessageContext
+from ._serialization import try_get_known_serializers_for_type
+from ._subscription import Subscription
+from ._subscription_context import SubscriptionInstantiationContext
+from ._topic import TopicId
+from ._type_helpers import get_types
+from .exceptions import CantHandleException
+
+T = TypeVar("T")
+ClosureAgentType = TypeVar("ClosureAgentType", bound="ClosureAgent")
+
+
+def get_handled_types_from_closure(
+    closure: Callable[[ClosureAgent, T, MessageContext], Awaitable[Any]],
+) -> Sequence[type]:
+    args = inspect.getfullargspec(closure)[0]
+    if len(args) != 3:
+        raise AssertionError("Closure must have 4 arguments")
+
+    message_arg_name = args[1]
+
+    type_hints = get_type_hints(closure)
+
+    if "return" not in type_hints:
+        raise AssertionError("return not found in function signature")
+
+    # Get the type of the message parameter
+    target_types = get_types(type_hints[message_arg_name])
+    if target_types is None:
+        raise AssertionError("Message type not found")
+
+    # print(type_hints)
+    return_types = get_types(type_hints["return"])
+
+    if return_types is None:
+        raise AssertionError("Return type not found")
+
+    return target_types
+
+
+class ClosureContext(Protocol):
+    @property
+    def id(self) -> AgentId: ...
+
+    async def send_message(
+        self,
+        message: Any,
+        recipient: AgentId,
+        *,
+        cancellation_token: CancellationToken | None = None,
+    ) -> Any: ...
+
+    async def publish_message(
+        self,
+        message: Any,
+        topic_id: TopicId,
+        *,
+        cancellation_token: CancellationToken | None = None,
+    ) -> None: ...
+
+
+class ClosureAgent(BaseAgent, ClosureContext):
+    def __init__(
+        self, description: str, closure: Callable[[ClosureContext, T, MessageContext], Awaitable[Any]]
+    ) -> None:
+        try:
+            runtime = AgentInstantiationContext.current_runtime()
+            id = AgentInstantiationContext.current_agent_id()
+        except Exception as e:
+            raise RuntimeError(
+                "ClosureAgent must be instantiated within the context of an AgentRuntime. It cannot be directly instantiated."
+            ) from e
+
+        self._runtime: AgentRuntime = runtime
+        self._id: AgentId = id
+        self._description = description
+        handled_types = get_handled_types_from_closure(closure)
+        self._expected_types = handled_types
+        self._closure = closure
+        super().__init__(description)
+
+    @property
+    def metadata(self) -> AgentMetadata:
+        assert self._id is not None
+        return AgentMetadata(
+            key=self._id.key,
+            type=self._id.type,
+            description=self._description,
+        )
+
+    @property
+    def id(self) -> AgentId:
+        return self._id
+
+    @property
+    def runtime(self) -> AgentRuntime:
+        return self._runtime
+
+    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}"
+            )
+        return await self._closure(self, message, ctx)
+
+    async def save_state(self) -> Mapping[str, Any]:
+        raise ValueError("save_state not implemented for ClosureAgent")
+
+    async def load_state(self, state: Mapping[str, Any]) -> None:
+        raise ValueError("load_state not implemented for ClosureAgent")
+
+    @classmethod
+    async def register_closure(
+        cls,
+        runtime: AgentRuntime,
+        type: str,
+        closure: Callable[[ClosureContext, T, MessageContext], Awaitable[Any]],
+        *,
+        skip_class_subscriptions: bool = False,
+        skip_direct_message_subscription: bool = False,
+        description: str = "",
+        subscriptions: Callable[[], list[Subscription] | Awaitable[list[Subscription]]] | None = None,
+    ) -> AgentType:
+        def factory() -> ClosureAgent:
+            return ClosureAgent(description=description, closure=closure)
+
+        agent_type = await cls.register(
+            runtime=runtime,
+            type=type,
+            factory=factory,  # type: ignore
+            skip_class_subscriptions=skip_class_subscriptions,
+            skip_direct_message_subscription=skip_direct_message_subscription,
+        )
+
+        subscriptions_list: List[Subscription] = []
+        if subscriptions is not None:
+            with SubscriptionInstantiationContext.populate_context(agent_type):
+                subscriptions_list_result = subscriptions()
+                if inspect.isawaitable(subscriptions_list_result):
+                    subscriptions_list.extend(await subscriptions_list_result)
+                else:
+                    # just ignore mypy here
+                    subscriptions_list.extend(subscriptions_list_result)  # type: ignore
+
+        for subscription in subscriptions_list:
+            await runtime.add_subscription(subscription)
+
+        handled_types = get_handled_types_from_closure(closure)
+        for message_type in handled_types:
+            # TODO: support custom serializers
+            serializer = try_get_known_serializers_for_type(message_type)
+            runtime.add_message_serializer(serializer)
+
+        return agent_type

python/packages/autogen-core/src/autogen_core/_default_subscription.py

@@ -1,8 +1,9 @@
 from typing import Callable, Type, TypeVar, overload
 
-from ..base import BaseAgent, SubscriptionInstantiationContext, subscription_factory
-from ..base.exceptions import CantHandleException
+from ._base_agent import BaseAgent, subscription_factory
+from ._subscription_context import SubscriptionInstantiationContext
 from ._type_subscription import TypeSubscription
+from .exceptions import CantHandleException
 
 
 class DefaultSubscription(TypeSubscription):

python/packages/autogen-core/src/autogen_core/_default_topic.py

@@ -1,4 +1,5 @@
-from ..base import MessageHandlerContext, TopicId
+from ._message_handler_context import MessageHandlerContext
+from ._topic import TopicId
 
 
 class DefaultTopicId(TopicId):

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

@@ -0,0 +1,518 @@
+import logging
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Coroutine,
+    DefaultDict,
+    List,
+    Literal,
+    Protocol,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    cast,
+    get_type_hints,
+    overload,
+    runtime_checkable,
+)
+
+from ._base_agent import BaseAgent
+from ._message_context import MessageContext
+from ._serialization import MessageSerializer, try_get_known_serializers_for_type
+from ._type_helpers import AnyType, get_types
+from .exceptions import CantHandleException
+
+logger = logging.getLogger("autogen_core")
+
+AgentT = TypeVar("AgentT")
+ReceivesT = TypeVar("ReceivesT")
+ProducesT = TypeVar("ProducesT", covariant=True)
+
+# TODO: Generic typevar bound binding U to agent type
+# Can't do because python doesnt support it
+
+
+# Pyright and mypy disagree on the variance of ReceivesT. Mypy thinks it should be contravariant here.
+# Revisit this later to see if we can remove the ignore.
+@runtime_checkable
+class MessageHandler(Protocol[AgentT, ReceivesT, ProducesT]):  # type: ignore
+    target_types: Sequence[type]
+    produces_types: Sequence[type]
+    is_message_handler: Literal[True]
+    router: Callable[[ReceivesT, MessageContext], bool]
+
+    # agent_instance binds to self in the method
+    @staticmethod
+    async def __call__(agent_instance: AgentT, message: ReceivesT, ctx: MessageContext) -> ProducesT: ...
+
+
+# NOTE: this works on concrete types and not inheritance
+# TODO: Use a protocol for the outer function to check checked arg names
+
+
+@overload
+def message_handler(
+    func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]],
+) -> MessageHandler[AgentT, ReceivesT, ProducesT]: ...
+
+
+@overload
+def message_handler(
+    func: None = None,
+    *,
+    match: None = ...,
+    strict: bool = ...,
+) -> Callable[
+    [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
+    MessageHandler[AgentT, ReceivesT, ProducesT],
+]: ...
+
+
+@overload
+def message_handler(
+    func: None = None,
+    *,
+    match: Callable[[ReceivesT, MessageContext], bool],
+    strict: bool = ...,
+) -> Callable[
+    [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
+    MessageHandler[AgentT, ReceivesT, ProducesT],
+]: ...
+
+
+def message_handler(
+    func: None | Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]] = None,
+    *,
+    strict: bool = True,
+    match: None | Callable[[ReceivesT, MessageContext], bool] = None,
+) -> (
+    Callable[
+        [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
+        MessageHandler[AgentT, ReceivesT, ProducesT],
+    ]
+    | MessageHandler[AgentT, ReceivesT, ProducesT]
+):
+    """Decorator for generic message handlers.
+
+    Add this decorator to methods in a :class:`RoutedAgent` class that are intended to handle both event and RPC messages.
+    These methods must have a specific signature that needs to be followed for it to be valid:
+
+    - The method must be an `async` method.
+    - The method must be decorated with the `@message_handler` decorator.
+    - The method must have exactly 3 arguments:
+        1. `self`
+        2. `message`: The message to be handled, this must be type-hinted with the message type that it is intended to handle.
+        3. `ctx`: A :class:`autogen_core.base.MessageContext` object.
+    - The method must be type hinted with what message types it can return as a response, or it can return `None` if it does not return anything.
+
+    Handlers can handle more than one message type by accepting a Union of the message types. It can also return more than one message type by returning a Union of the message types.
+
+    Args:
+        func: The function to be decorated.
+        strict: If `True`, the handler will raise an exception if the message type or return type is not in the target types. If `False`, it will log a warning instead.
+        match: A function that takes the message and the context as arguments and returns a boolean. This is used for secondary routing after the message type. For handlers addressing the same message type, the match function is applied in alphabetical order of the handlers and the first matching handler will be called while the rest are skipped. If `None`, the first handler in alphabetical order matching the same message type will be called.
+    """
+
+    def decorator(
+        func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]],
+    ) -> MessageHandler[AgentT, ReceivesT, ProducesT]:
+        type_hints = get_type_hints(func)
+        if "message" not in type_hints:
+            raise AssertionError("message parameter not found in function signature")
+
+        if "return" not in type_hints:
+            raise AssertionError("return not found in function signature")
+
+        # Get the type of the message parameter
+        target_types = get_types(type_hints["message"])
+        if target_types is None:
+            raise AssertionError("Message type not found")
+
+        # print(type_hints)
+        return_types = get_types(type_hints["return"])
+
+        if return_types is None:
+            raise AssertionError("Return type not found")
+
+        # Convert target_types to list and stash
+
+        @wraps(func)
+        async def wrapper(self: AgentT, message: ReceivesT, ctx: MessageContext) -> ProducesT:
+            if type(message) not in target_types:
+                if strict:
+                    raise CantHandleException(f"Message type {type(message)} not in target types {target_types}")
+                else:
+                    logger.warning(f"Message type {type(message)} not in target types {target_types}")
+
+            return_value = await func(self, message, ctx)
+
+            if AnyType not in return_types and type(return_value) not in return_types:
+                if strict:
+                    raise ValueError(f"Return type {type(return_value)} not in return types {return_types}")
+                else:
+                    logger.warning(f"Return type {type(return_value)} not in return types {return_types}")
+
+            return return_value
+
+        wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, ProducesT], wrapper)
+        wrapper_handler.target_types = list(target_types)
+        wrapper_handler.produces_types = list(return_types)
+        wrapper_handler.is_message_handler = True
+        wrapper_handler.router = match or (lambda _message, _ctx: True)
+
+        return wrapper_handler
+
+    if func is None and not callable(func):
+        return decorator
+    elif callable(func):
+        return decorator(func)
+    else:
+        raise ValueError("Invalid arguments")
+
+
+@overload
+def event(
+    func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]],
+) -> MessageHandler[AgentT, ReceivesT, None]: ...
+
+
+@overload
+def event(
+    func: None = None,
+    *,
+    match: None = ...,
+    strict: bool = ...,
+) -> Callable[
+    [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]]],
+    MessageHandler[AgentT, ReceivesT, None],
+]: ...
+
+
+@overload
+def event(
+    func: None = None,
+    *,
+    match: Callable[[ReceivesT, MessageContext], bool],
+    strict: bool = ...,
+) -> Callable[
+    [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]]],
+    MessageHandler[AgentT, ReceivesT, None],
+]: ...
+
+
+def event(
+    func: None | Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]] = None,
+    *,
+    strict: bool = True,
+    match: None | Callable[[ReceivesT, MessageContext], bool] = None,
+) -> (
+    Callable[
+        [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]]],
+        MessageHandler[AgentT, ReceivesT, None],
+    ]
+    | MessageHandler[AgentT, ReceivesT, None]
+):
+    """Decorator for event message handlers.
+
+    Add this decorator to methods in a :class:`RoutedAgent` class that are intended to handle event messages.
+    These methods must have a specific signature that needs to be followed for it to be valid:
+
+    - The method must be an `async` method.
+    - The method must be decorated with the `@message_handler` decorator.
+    - The method must have exactly 3 arguments:
+        1. `self`
+        2. `message`: The event message to be handled, this must be type-hinted with the message type that it is intended to handle.
+        3. `ctx`: A :class:`autogen_core.base.MessageContext` object.
+    - The method must return `None`.
+
+    Handlers can handle more than one message type by accepting a Union of the message types.
+
+    Args:
+        func: The function to be decorated.
+        strict: If `True`, the handler will raise an exception if the message type is not in the target types. If `False`, it will log a warning instead.
+        match: A function that takes the message and the context as arguments and returns a boolean. This is used for secondary routing after the message type. For handlers addressing the same message type, the match function is applied in alphabetical order of the handlers and the first matching handler will be called while the rest are skipped. If `None`, the first handler in alphabetical order matching the same message type will be called.
+    """
+
+    def decorator(
+        func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]],
+    ) -> MessageHandler[AgentT, ReceivesT, None]:
+        type_hints = get_type_hints(func)
+        if "message" not in type_hints:
+            raise AssertionError("message parameter not found in function signature")
+
+        if "return" not in type_hints:
+            raise AssertionError("return not found in function signature")
+
+        # Get the type of the message parameter
+        target_types = get_types(type_hints["message"])
+        if target_types is None:
+            raise AssertionError("Message type not found. Please provide a type hint for the message parameter.")
+
+        return_types = get_types(type_hints["return"])
+
+        if return_types is None:
+            raise AssertionError("Return type not found. Please use `None` as the type hint of the return type.")
+
+        # Convert target_types to list and stash
+
+        @wraps(func)
+        async def wrapper(self: AgentT, message: ReceivesT, ctx: MessageContext) -> None:
+            if type(message) not in target_types:
+                if strict:
+                    raise CantHandleException(f"Message type {type(message)} not in target types {target_types}")
+                else:
+                    logger.warning(f"Message type {type(message)} not in target types {target_types}")
+
+            return_value = await func(self, message, ctx)  # type: ignore
+
+            if return_value is not None:
+                if strict:
+                    raise ValueError(f"Return type {type(return_value)} is not None.")
+                else:
+                    logger.warning(f"Return type {type(return_value)} is not None. It will be ignored.")
+
+            return None
+
+        wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, None], wrapper)
+        wrapper_handler.target_types = list(target_types)
+        wrapper_handler.produces_types = list(return_types)
+        wrapper_handler.is_message_handler = True
+        # Wrap the match function with a check on the is_rpc flag.
+        wrapper_handler.router = lambda _message, _ctx: (not _ctx.is_rpc) and (match(_message, _ctx) if match else True)
+
+        return wrapper_handler
+
+    if func is None and not callable(func):
+        return decorator
+    elif callable(func):
+        return decorator(func)
+    else:
+        raise ValueError("Invalid arguments")
+
+
+@overload
+def rpc(
+    func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]],
+) -> MessageHandler[AgentT, ReceivesT, ProducesT]: ...
+
+
+@overload
+def rpc(
+    func: None = None,
+    *,
+    match: None = ...,
+    strict: bool = ...,
+) -> Callable[
+    [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
+    MessageHandler[AgentT, ReceivesT, ProducesT],
+]: ...
+
+
+@overload
+def rpc(
+    func: None = None,
+    *,
+    match: Callable[[ReceivesT, MessageContext], bool],
+    strict: bool = ...,
+) -> Callable[
+    [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
+    MessageHandler[AgentT, ReceivesT, ProducesT],
+]: ...
+
+
+def rpc(
+    func: None | Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]] = None,
+    *,
+    strict: bool = True,
+    match: None | Callable[[ReceivesT, MessageContext], bool] = None,
+) -> (
+    Callable[
+        [Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
+        MessageHandler[AgentT, ReceivesT, ProducesT],
+    ]
+    | MessageHandler[AgentT, ReceivesT, ProducesT]
+):
+    """Decorator for RPC message handlers.
+
+    Add this decorator to methods in a :class:`RoutedAgent` class that are intended to handle RPC messages.
+    These methods must have a specific signature that needs to be followed for it to be valid:
+
+    - The method must be an `async` method.
+    - The method must be decorated with the `@message_handler` decorator.
+    - The method must have exactly 3 arguments:
+        1. `self`
+        2. `message`: The message to be handled, this must be type-hinted with the message type that it is intended to handle.
+        3. `ctx`: A :class:`autogen_core.base.MessageContext` object.
+    - The method must be type hinted with what message types it can return as a response, or it can return `None` if it does not return anything.
+
+    Handlers can handle more than one message type by accepting a Union of the message types. It can also return more than one message type by returning a Union of the message types.
+
+    Args:
+        func: The function to be decorated.
+        strict: If `True`, the handler will raise an exception if the message type or return type is not in the target types. If `False`, it will log a warning instead.
+        match: A function that takes the message and the context as arguments and returns a boolean. This is used for secondary routing after the message type. For handlers addressing the same message type, the match function is applied in alphabetical order of the handlers and the first matching handler will be called while the rest are skipped. If `None`, the first handler in alphabetical order matching the same message type will be called.
+    """
+
+    def decorator(
+        func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]],
+    ) -> MessageHandler[AgentT, ReceivesT, ProducesT]:
+        type_hints = get_type_hints(func)
+        if "message" not in type_hints:
+            raise AssertionError("message parameter not found in function signature")
+
+        if "return" not in type_hints:
+            raise AssertionError("return not found in function signature")
+
+        # Get the type of the message parameter
+        target_types = get_types(type_hints["message"])
+        if target_types is None:
+            raise AssertionError("Message type not found")
+
+        # print(type_hints)
+        return_types = get_types(type_hints["return"])
+
+        if return_types is None:
+            raise AssertionError("Return type not found")
+
+        # Convert target_types to list and stash
+
+        @wraps(func)
+        async def wrapper(self: AgentT, message: ReceivesT, ctx: MessageContext) -> ProducesT:
+            if type(message) not in target_types:
+                if strict:
+                    raise CantHandleException(f"Message type {type(message)} not in target types {target_types}")
+                else:
+                    logger.warning(f"Message type {type(message)} not in target types {target_types}")
+
+            return_value = await func(self, message, ctx)
+
+            if AnyType not in return_types and type(return_value) not in return_types:
+                if strict:
+                    raise ValueError(f"Return type {type(return_value)} not in return types {return_types}")
+                else:
+                    logger.warning(f"Return type {type(return_value)} not in return types {return_types}")
+
+            return return_value
+
+        wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, ProducesT], wrapper)
+        wrapper_handler.target_types = list(target_types)
+        wrapper_handler.produces_types = list(return_types)
+        wrapper_handler.is_message_handler = True
+        wrapper_handler.router = lambda _message, _ctx: (_ctx.is_rpc) and (match(_message, _ctx) if match else True)
+
+        return wrapper_handler
+
+    if func is None and not callable(func):
+        return decorator
+    elif callable(func):
+        return decorator(func)
+    else:
+        raise ValueError("Invalid arguments")
+
+
+class RoutedAgent(BaseAgent):
+    """A base class for agents that route messages to handlers based on the type of the message
+    and optional matching functions.
+
+    To create a routed agent, subclass this class and add message handlers as methods decorated with
+    either :func:`event` or :func:`rpc` decorator.
+
+    Example:
+
+    .. code-block:: python
+
+        from dataclasses import dataclass
+        from autogen_core import MessageContext
+        from autogen_core import RoutedAgent, event, rpc
+
+
+        @dataclass
+        class Message:
+            pass
+
+
+        @dataclass
+        class MessageWithContent:
+            content: str
+
+
+        @dataclass
+        class Response:
+            pass
+
+
+        class MyAgent(RoutedAgent):
+            def __init__(self):
+                super().__init__("MyAgent")
+
+            @event
+            async def handle_event_message(self, message: Message, ctx: MessageContext) -> None:
+                assert ctx.topic_id is not None
+                await self.publish_message(MessageWithContent("event handled"), ctx.topic_id)
+
+            @rpc(match=lambda message, ctx: message.content == "special")  # type: ignore
+            async def handle_special_rpc_message(self, message: MessageWithContent, ctx: MessageContext) -> Response:
+                return Response()
+    """
+
+    def __init__(self, description: str) -> None:
+        # Self is already bound to the handlers
+        self._handlers: DefaultDict[
+            Type[Any],
+            List[MessageHandler[RoutedAgent, Any, Any]],
+        ] = DefaultDict(list)
+
+        handlers = self._discover_handlers()
+        for message_handler in handlers:
+            for target_type in message_handler.target_types:
+                self._handlers[target_type].append(message_handler)
+
+        super().__init__(description)
+
+    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."""
+        key_type: Type[Any] = type(message)  # type: ignore
+        handlers = self._handlers.get(key_type)  # type: ignore
+        if handlers is not None:
+            # Iterate over all handlers for this matching message type.
+            # Call the first handler whose router returns True and then return the result.
+            for h in handlers:
+                if h.router(message, ctx):
+                    return await h(self, message, ctx)
+        return await self.on_unhandled_message(message, ctx)  # type: ignore
+
+    async def on_unhandled_message(self, message: Any, ctx: MessageContext) -> None:
+        """Called when a message is received that does not have a matching message handler.
+        The default implementation logs an info message."""
+        logger.info(f"Unhandled message: {message}")
+
+    @classmethod
+    def _discover_handlers(cls) -> Sequence[MessageHandler[Any, Any, Any]]:
+        handlers: List[MessageHandler[Any, Any, Any]] = []
+        for attr in dir(cls):
+            if callable(getattr(cls, attr, None)):
+                # Since we are getting it from the class, self is not bound
+                handler = getattr(cls, attr)
+                if hasattr(handler, "is_message_handler"):
+                    handlers.append(cast(MessageHandler[Any, Any, Any], handler))
+        return handlers
+
+    @classmethod
+    def _handles_types(cls) -> List[Tuple[Type[Any], List[MessageSerializer[Any]]]]:
+        # TODO handle deduplication
+        handlers = cls._discover_handlers()
+        types: List[Tuple[Type[Any], List[MessageSerializer[Any]]]] = []
+        types.extend(cls.internal_extra_handles_types)
+        for handler in handlers:
+            for t in handler.target_types:
+                # TODO: support different serializers
+                serializers = try_get_known_serializers_for_type(t)
+                if len(serializers) == 0:
+                    raise ValueError(f"No serializers found for type {t}.")
+
+                types.append((t, try_get_known_serializers_for_type(t)))
+        return types

python/packages/autogen-core/src/autogen_core/_serialization.py

@@ -6,7 +6,7 @@ from google.protobuf import any_pb2
 from google.protobuf.message import Message
 from pydantic import BaseModel
 
-from autogen_core.base._type_helpers import is_union
+from ._type_helpers import is_union
 
 T = TypeVar("T")
 

python/packages/autogen-core/src/autogen_core/_subscription_context.py

@@ -2,7 +2,7 @@ from contextlib import contextmanager
 from contextvars import ContextVar
 from typing import Any, ClassVar, Generator
 
-from autogen_core.base._agent_type import AgentType
+from autogen_core._agent_type import AgentType
 
 
 class SubscriptionInstantiationContext:

python/packages/autogen-core/src/autogen_core/_type_prefix_subscription.py

@@ -15,7 +15,7 @@ class TypePrefixSubscription(Subscription):
 
         .. code-block:: python
 
-            from autogen_core.components import TypePrefixSubscription
+            from autogen_core import TypePrefixSubscription
 
             subscription = TypePrefixSubscription(topic_type_prefix="t1", agent_type="a1")
 

python/packages/autogen-core/src/autogen_core/_type_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 TypeSubscription(Subscription):
@@ -13,7 +15,7 @@ class TypeSubscription(Subscription):
 
         .. code-block:: python
 
-            from autogen_core.components import TypeSubscription
+            from autogen_core import TypeSubscription
 
             subscription = TypeSubscription(topic_type="t1", agent_type="a1")
 

python/packages/autogen-core/src/autogen_core/application/_helpers.py

@@ -1,11 +1,11 @@
 from collections import defaultdict
 from typing import Awaitable, Callable, DefaultDict, List, Set
 
-from ..base._agent import Agent
-from ..base._agent_id import AgentId
-from ..base._agent_type import AgentType
-from ..base._subscription import Subscription
-from ..base._topic import TopicId
+from .._agent import Agent
+from .._agent_id import AgentId
+from .._agent_type import AgentType
+from .._subscription import Subscription
+from .._topic import TopicId
 
 
 async def get_impl(

python/packages/autogen-core/src/autogen_core/application/_single_threaded_agent_runtime.py

@@ -15,9 +15,9 @@ from typing import Any, Awaitable, Callable, Dict, List, Mapping, ParamSpec, Set
 from opentelemetry.trace import TracerProvider
 from typing_extensions import deprecated
 
-from autogen_core.base._serialization import MessageSerializer, SerializationRegistry
+from autogen_core._serialization import MessageSerializer, SerializationRegistry
 
-from ..base import (
+from .. import (
     Agent,
     AgentId,
     AgentInstantiationContext,
@@ -31,8 +31,8 @@ from ..base import (
     SubscriptionInstantiationContext,
     TopicId,
 )
-from ..base.exceptions import MessageDroppedException
 from ..base.intervention import DropMessage, InterventionHandler
+from ..exceptions import MessageDroppedException
 from ._helpers import SubscriptionManager, get_impl
 from .telemetry import EnvelopeMetadata, MessageRuntimeTracingConfig, TraceHelper, get_telemetry_envelope_metadata
 

python/packages/autogen-core/src/autogen_core/application/_worker_runtime.py

@@ -34,9 +34,7 @@ from typing_extensions import Self, deprecated
 
 from autogen_core.application.protos import cloudevent_pb2
 
-from ..base import (
-    JSON_DATA_CONTENT_TYPE,
-    PROTOBUF_DATA_CONTENT_TYPE,
+from .. import (
     Agent,
     AgentId,
     AgentInstantiationContext,
@@ -50,9 +48,15 @@ from ..base import (
     SubscriptionInstantiationContext,
     TopicId,
 )
-from ..base._serialization import MessageSerializer, SerializationRegistry
-from ..base._type_helpers import ChannelArgumentType
-from ..components import TypePrefixSubscription, TypeSubscription
+from .._serialization import (
+    JSON_DATA_CONTENT_TYPE,
+    PROTOBUF_DATA_CONTENT_TYPE,
+    MessageSerializer,
+    SerializationRegistry,
+)
+from .._type_helpers import ChannelArgumentType
+from .._type_prefix_subscription import TypePrefixSubscription
+from .._type_subscription import TypeSubscription
 from . import _constants
 from ._constants import GRPC_IMPORT_ERROR_STR
 from ._helpers import SubscriptionManager, get_impl

python/packages/autogen-core/src/autogen_core/application/_worker_runtime_host.py

@@ -3,7 +3,7 @@ import logging
 import signal
 from typing import Optional, Sequence
 
-from ..base._type_helpers import ChannelArgumentType
+from .._type_helpers import ChannelArgumentType
 from ._constants import GRPC_IMPORT_ERROR_STR
 from ._worker_runtime_host_servicer import WorkerAgentRuntimeHostServicer
 

python/packages/autogen-core/src/autogen_core/application/_worker_runtime_host_servicer.py

@@ -4,10 +4,8 @@ from _collections_abc import AsyncIterator, Iterator
 from asyncio import Future, Task
 from typing import Any, Dict, Set, cast
 
-from autogen_core.base._type_prefix_subscription import TypePrefixSubscription
-
-from ..base import Subscription, TopicId
-from ..components import TypeSubscription
+from .. import Subscription, TopicId, TypeSubscription
+from .._type_prefix_subscription import TypePrefixSubscription
 from ._constants import GRPC_IMPORT_ERROR_STR
 from ._helpers import SubscriptionManager