Flatten core base and components (#4513)

* Flatten core base and components

* remove extra files

* dont export from deprecated locations

* format

* fmt

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,

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,

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 import CancellationToken, FunctionCall
-from autogen_core.components import 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.ui 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 import CancellationToken, Image
-from autogen_core.components import 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/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 import (
-from autogen_core.base import (
     AgentId,
     AgentInstantiationContext,
     AgentRuntime,
     AgentType,
     CancellationToken,
+    ClosureAgent,
     MessageContext,
+    TypeSubscription,
 )
-from autogen_core.components import ClosureAgent, TypeSubscription
+from autogen_core._closure_agent import ClosureContext
-from autogen_core.components._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.conditions 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
 
 
@@ -317,7 +318,7 @@ class BaseGroupChat(Team, ABC):
             from autogen_agentchat.conditions import MaxMessageTermination
             from autogen_agentchat.ui import 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 import DefaultTopicId, MessageContext, event, rpc
-from autogen_core.components import DefaultTopicId, 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 import DefaultTopicId, MessageContext, event, rpc
-from autogen_core.components import DefaultTopicId, 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 import AgentId, CancellationToken, DefaultTopicId, Image, MessageContext, event, rpc
-from autogen_core.components import DefaultTopicId, Image, 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 import MessageContext, RoutedAgent
-from autogen_core.components import RoutedAgent
 
 
 class FIFOLock:

python/packages/autogen-agentchat/src/autogen_agentchat/ui/_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/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

@@ -29,8 +29,7 @@ from autogen_agentchat.teams import (
     Swarm,
 )
 from autogen_agentchat.ui import Console
-from autogen_core.base import CancellationToken
+from autogen_core import CancellationToken, FunctionCall
-from autogen_core.components import 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 import ClosureAgent, ClosureContext, DefaultSubscription, DefaultTopicId, MessageContext\n",
-    "from autogen_core.base import MessageContext\n",
+    "from autogen_core.application import SingleThreadedAgentRuntime"
-    "from autogen_core.components import ClosureAgent, ClosureContext, DefaultSubscription, DefaultTopicId"
    ]
   },
   {

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 import AgentId, MessageContext, RoutedAgent, message_handler\n",
-    "from autogen_core.components import 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"
-    "from autogen_core.base.intervention import DefaultInterventionHandler\n",
-    "from autogen_core.components import DefaultTopicId, RoutedAgent, default_subscription, message_handler"
    ]
   },
   {

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

@@ -3,23 +3,20 @@ from __future__ import annotations
 import inspect
 from typing import Any, Awaitable, Callable, List, Mapping, Protocol, Sequence, TypeVar, get_type_hints
 
-from autogen_core.base._serialization import try_get_known_serializers_for_type
+from ._agent_id import AgentId
-from autogen_core.base._subscription_context import SubscriptionInstantiationContext
+from ._agent_instantiation import AgentInstantiationContext
-
+from ._agent_metadata import AgentMetadata
-from ..base import (
+from ._agent_runtime import AgentRuntime
-    AgentId,
+from ._agent_type import AgentType
-    AgentInstantiationContext,
+from ._base_agent import BaseAgent
-    AgentMetadata,
+from ._cancellation_token import CancellationToken
-    AgentRuntime,
+from ._message_context import MessageContext
-    AgentType,
+from ._serialization import try_get_known_serializers_for_type
-    BaseAgent,
+from ._subscription import Subscription
-    CancellationToken,
+from ._subscription_context import SubscriptionInstantiationContext
-    MessageContext,
+from ._topic import TopicId
-    Subscription,
+from ._type_helpers import get_types
-    TopicId,
+from .exceptions import CantHandleException
-)
-from ..base._type_helpers import get_types
-from ..base.exceptions import CantHandleException
 
 T = TypeVar("T")
 ClosureAgentType = TypeVar("ClosureAgentType", bound="ClosureAgent")

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_agent import BaseAgent, subscription_factory
-from ..base.exceptions import CantHandleException
+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

@@ -1,5 +1,4 @@
 import logging
-import warnings
 from functools import wraps
 from typing import (
     Any,
@@ -19,9 +18,11 @@ from typing import (
     runtime_checkable,
 )
 
-from ..base import BaseAgent, MessageContext, MessageSerializer, try_get_known_serializers_for_type
+from ._base_agent import BaseAgent
-from ..base._type_helpers import AnyType, get_types
+from ._message_context import MessageContext
-from ..base.exceptions import CantHandleException
+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")
 
@@ -423,8 +424,8 @@ class RoutedAgent(BaseAgent):
     .. code-block:: python
 
         from dataclasses import dataclass
-        from autogen_core.base import MessageContext
+        from autogen_core import MessageContext
-        from autogen_core.components import RoutedAgent, event, rpc
+        from autogen_core import RoutedAgent, event, rpc
 
 
         @dataclass
@@ -515,12 +516,3 @@ class RoutedAgent(BaseAgent):
 
                 types.append((t, try_get_known_serializers_for_type(t)))
         return types
-
-
-# Deprecation warning for TypeRoutedAgent
-class TypeRoutedAgent(RoutedAgent):
-    """Deprecated. Use :class:`RoutedAgent` instead."""
-
-    def __init__(self, description: str) -> None:
-        warnings.warn("TypeRoutedAgent is deprecated. Use RoutedAgent instead.", DeprecationWarning, stacklevel=2)
-        super().__init__(description)

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 ._agent_id import AgentId
-from ..base.exceptions import CantHandleException
+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 .._agent import Agent
-from ..base._agent_id import AgentId
+from .._agent_id import AgentId
-from ..base._agent_type import AgentType
+from .._agent_type import AgentType
-from ..base._subscription import Subscription
+from .._subscription import Subscription
-from ..base._topic import TopicId
+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 (
+from .. import (
-    JSON_DATA_CONTENT_TYPE,
-    PROTOBUF_DATA_CONTENT_TYPE,
     Agent,
     AgentId,
     AgentInstantiationContext,
@@ -50,9 +48,15 @@ from ..base import (
     SubscriptionInstantiationContext,
     TopicId,
 )
-from ..base._serialization import MessageSerializer, SerializationRegistry
+from .._serialization import (
-from ..base._type_helpers import ChannelArgumentType
+    JSON_DATA_CONTENT_TYPE,
-from ..components import TypePrefixSubscription, TypeSubscription
+    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 .. import Subscription, TopicId, TypeSubscription
-
+from .._type_prefix_subscription import TypePrefixSubscription
-from ..base import Subscription, TopicId
-from ..components import TypeSubscription
 from ._constants import GRPC_IMPORT_ERROR_STR
 from ._helpers import SubscriptionManager