Merge branch 'main' into distributed_agent_startup

Author: ekzhu

.github/workflows/docs.yml

@@ -33,7 +33,7 @@ jobs:
           [
             # For main use the workflow target
             { ref: "${{github.ref}}", dest-dir: dev, uv-version: "0.5.13", sphinx-release-override: "dev" },
-            { ref: "v0.4.2", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" },
+            { ref: "v0.4.3", dest-dir: stable, uv-version: "0.5.13", sphinx-release-override: "stable" },
             { ref: "v0.4.0.dev0", dest-dir: "0.4.0.dev0", uv-version: "0.5.11", sphinx-release-override: "" },
             { ref: "v0.4.0.dev1", dest-dir: "0.4.0.dev1", uv-version: "0.5.11", sphinx-release-override: "" },
             { ref: "v0.4.0.dev2", dest-dir: "0.4.0.dev2", uv-version: "0.5.11", sphinx-release-override: "" },
@@ -51,6 +51,7 @@ jobs:
             { ref: "v0.4.0.post1", dest-dir: "0.4.0", uv-version: "0.5.13", sphinx-release-override: "" },
             { ref: "v0.4.1", dest-dir: "0.4.1", uv-version: "0.5.13", sphinx-release-override: "" },
             { ref: "v0.4.2", dest-dir: "0.4.2", uv-version: "0.5.13", sphinx-release-override: "" },
+            { ref: "v0.4.3", dest-dir: "0.4.3", uv-version: "0.5.13", sphinx-release-override: "" },
           ]
     steps:
       - name: Checkout

docs/switcher.json

@@ -1,6 +1,6 @@
 [
     {
-        "name": "0.4.2 (stable)",
+        "name": "0.4.3 (stable)",
         "version": "stable",
         "url": "/autogen/stable/",
         "preferred": true

protos/agent_worker.proto

@@ -48,12 +48,12 @@ message Event {
 }
 
 message RegisterAgentTypeRequest {
-    string request_id = 1;
+    string request_id = 1; // TODO: remove once message based requests are removed
     string type = 2;
 }
 
 message RegisterAgentTypeResponse {
-    string request_id = 1;
+    string request_id = 1; // TODO: remove once message based requests are removed
     bool success = 2;
     optional string error = 3;
 }
@@ -69,27 +69,46 @@ message TypePrefixSubscription {
 }
 
 message Subscription {
+    string id = 1;
     oneof subscription {
-        TypeSubscription typeSubscription = 1;
-        TypePrefixSubscription typePrefixSubscription = 2;
+        TypeSubscription typeSubscription = 2;
+        TypePrefixSubscription typePrefixSubscription = 3;
     }
 }
 
 message AddSubscriptionRequest {
-    string request_id = 1;
+    string request_id = 1; // TODO: remove once message based requests are removed
     Subscription subscription = 2;
 }
 
 message AddSubscriptionResponse {
-    string request_id = 1;
+    string request_id = 1; // TODO: remove once message based requests are removed
     bool success = 2;
     optional string error = 3;
 }
 
+message RemoveSubscriptionRequest {
+    string id = 1;
+}
+
+message RemoveSubscriptionResponse {
+    bool success = 1;
+    optional string error = 2;
+}
+
+message GetSubscriptionsRequest {}
+message GetSubscriptionsResponse {
+    repeated Subscription subscriptions = 1;
+}
+
 service AgentRpc {
     rpc OpenChannel (stream Message) returns (stream Message);
     rpc GetState(AgentId) returns (GetStateResponse);
     rpc SaveState(AgentState) returns (SaveStateResponse);
+    rpc RegisterAgent(RegisterAgentTypeRequest) returns (RegisterAgentTypeResponse);
+    rpc AddSubscription(AddSubscriptionRequest) returns (AddSubscriptionResponse);
+    rpc RemoveSubscription(RemoveSubscriptionRequest) returns (RemoveSubscriptionResponse);
+    rpc GetSubscriptions(GetSubscriptionsRequest) returns (GetSubscriptionsResponse);
 }
 
 message AgentState {

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

@@ -1,7 +1,9 @@
 from typing import Any, AsyncGenerator, List, Mapping, Sequence
 
-from autogen_core import CancellationToken
+from autogen_core import CancellationToken, Component, ComponentModel
 from autogen_core.models import ChatCompletionClient, LLMMessage, SystemMessage, UserMessage
+from pydantic import BaseModel
+from typing_extensions import Self
 
 from autogen_agentchat.base import Response
 from autogen_agentchat.state import SocietyOfMindAgentState
@@ -16,7 +18,18 @@
 from ._base_chat_agent import BaseChatAgent
 
 
-class SocietyOfMindAgent(BaseChatAgent):
+class SocietyOfMindAgentConfig(BaseModel):
+    """The declarative configuration for a SocietyOfMindAgent."""
+
+    name: str
+    team: ComponentModel
+    model_client: ComponentModel
+    description: str
+    instruction: str
+    response_prompt: str
+
+
+class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
     """An agent that uses an inner team of agents to generate responses.
 
     Each time the agent's :meth:`on_messages` or :meth:`on_messages_stream`
@@ -74,6 +87,9 @@ async def main() -> None:
         asyncio.run(main())
     """
 
+    component_config_schema = SocietyOfMindAgentConfig
+    component_provider_override = "autogen_agentchat.agents.SocietyOfMindAgent"
+
     DEFAULT_INSTRUCTION = "Earlier you were asked to fulfill a request. You and your team worked diligently to address that request. Here is a transcript of that conversation:"
     """str: The default instruction to use when generating a response using the
     inner team's messages. The instruction will be prepended to the inner team's
@@ -173,3 +189,26 @@ async def save_state(self) -> Mapping[str, Any]:
     async def load_state(self, state: Mapping[str, Any]) -> None:
         society_of_mind_state = SocietyOfMindAgentState.model_validate(state)
         await self._team.load_state(society_of_mind_state.inner_team_state)
+
+    def _to_config(self) -> SocietyOfMindAgentConfig:
+        return SocietyOfMindAgentConfig(
+            name=self.name,
+            team=self._team.dump_component(),
+            model_client=self._model_client.dump_component(),
+            description=self.description,
+            instruction=self._instruction,
+            response_prompt=self._response_prompt,
+        )
+
+    @classmethod
+    def _from_config(cls, config: SocietyOfMindAgentConfig) -> Self:
+        model_client = ChatCompletionClient.load_component(config.model_client)
+        team = Team.load_component(config.team)
+        return cls(
+            name=config.name,
+            team=team,
+            model_client=model_client,
+            description=config.description,
+            instruction=config.instruction,
+            response_prompt=config.response_prompt,
+        )

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

@@ -2,7 +2,8 @@
 from dataclasses import dataclass
 from typing import Any, AsyncGenerator, Mapping, Sequence
 
-from autogen_core import CancellationToken
+from autogen_core import CancellationToken, ComponentBase
+from pydantic import BaseModel
 
 from ..messages import AgentEvent, ChatMessage
 from ._task import TaskRunner
@@ -20,9 +21,11 @@ class Response:
     or :class:`ChatMessage`."""
 
 
-class ChatAgent(ABC, TaskRunner):
+class ChatAgent(ABC, TaskRunner, ComponentBase[BaseModel]):
     """Protocol for a chat agent."""
 
+    component_type = "agent"
+
     @property
     @abstractmethod
     def name(self) -> str:

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

@@ -1,17 +1,26 @@
-from typing import Any, Mapping, Protocol
+from abc import ABC, abstractmethod
+from typing import Any, Mapping
+
+from autogen_core import ComponentBase
+from pydantic import BaseModel
 
 from ._task import TaskRunner
 
 
-class Team(TaskRunner, Protocol):
+class Team(ABC, TaskRunner, ComponentBase[BaseModel]):
+    component_type = "team"
+
+    @abstractmethod
     async def reset(self) -> None:
         """Reset the team and all its participants to its initial state."""
         ...
 
+    @abstractmethod
     async def save_state(self) -> Mapping[str, Any]:
         """Save the current state of the team."""
         ...
 
+    @abstractmethod
     async def load_state(self, state: Mapping[str, Any]) -> None:
         """Load the state of the team."""
         ...

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

@@ -11,11 +11,13 @@
     AgentType,
     CancellationToken,
     ClosureAgent,
+    ComponentBase,
     MessageContext,
     SingleThreadedAgentRuntime,
     TypeSubscription,
 )
 from autogen_core._closure_agent import ClosureContext
+from pydantic import BaseModel
 
 from ... import EVENT_LOGGER_NAME
 from ...base import ChatAgent, TaskResult, Team, TerminationCondition
@@ -28,13 +30,15 @@
 event_logger = logging.getLogger(EVENT_LOGGER_NAME)
 
 
-class BaseGroupChat(Team, ABC):
+class BaseGroupChat(Team, ABC, ComponentBase[BaseModel]):
     """The base class for group chat teams.
 
     To implement a group chat team, first create a subclass of :class:`BaseGroupChatManager` and then
     create a subclass of :class:`BaseGroupChat` that uses the group chat manager.
     """
 
+    component_type = "team"
+
     def __init__(
         self,
         participants: List[ChatAgent],

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

@@ -1,7 +1,10 @@
 import logging
 from typing import Callable, List
 
+from autogen_core import Component, ComponentModel
 from autogen_core.models import ChatCompletionClient
+from pydantic import BaseModel
+from typing_extensions import Self
 
 from .... import EVENT_LOGGER_NAME, TRACE_LOGGER_NAME
 from ....base import ChatAgent, TerminationCondition
@@ -13,7 +16,18 @@
 event_logger = logging.getLogger(EVENT_LOGGER_NAME)
 
 
-class MagenticOneGroupChat(BaseGroupChat):
+class MagenticOneGroupChatConfig(BaseModel):
+    """The declarative configuration for a MagenticOneGroupChat."""
+
+    participants: List[ComponentModel]
+    model_client: ComponentModel
+    termination_condition: ComponentModel | None = None
+    max_turns: int | None = None
+    max_stalls: int
+    final_answer_prompt: str
+
+
+class MagenticOneGroupChat(BaseGroupChat, Component[MagenticOneGroupChatConfig]):
     """A team that runs a group chat with participants managed by the MagenticOneOrchestrator.
 
     The orchestrator handles the conversation flow, ensuring that the task is completed
@@ -73,6 +87,9 @@ async def main() -> None:
             }
     """
 
+    component_config_schema = MagenticOneGroupChatConfig
+    component_provider_override = "autogen_agentchat.teams.MagenticOneGroupChat"
+
     def __init__(
         self,
         participants: List[ChatAgent],
@@ -117,3 +134,31 @@ def _create_group_chat_manager_factory(
             self._final_answer_prompt,
             termination_condition,
         )
+
+    def _to_config(self) -> MagenticOneGroupChatConfig:
+        participants = [participant.dump_component() for participant in self._participants]
+        termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
+        return MagenticOneGroupChatConfig(
+            participants=participants,
+            model_client=self._model_client.dump_component(),
+            termination_condition=termination_condition,
+            max_turns=self._max_turns,
+            max_stalls=self._max_stalls,
+            final_answer_prompt=self._final_answer_prompt,
+        )
+
+    @classmethod
+    def _from_config(cls, config: MagenticOneGroupChatConfig) -> Self:
+        participants = [ChatAgent.load_component(participant) for participant in config.participants]
+        model_client = ChatCompletionClient.load_component(config.model_client)
+        termination_condition = (
+            TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
+        )
+        return cls(
+            participants,
+            model_client,
+            termination_condition=termination_condition,
+            max_turns=config.max_turns,
+            max_stalls=config.max_stalls,
+            final_answer_prompt=config.final_answer_prompt,
+        )

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

@@ -460,7 +460,7 @@ def _thread_to_context(self) -> List[LLMMessage]:
                 assert isinstance(m, TextMessage | ToolCallSummaryMessage)
                 context.append(AssistantMessage(content=m.content, source=m.source))
             else:
-                assert isinstance(m, TextMessage) or isinstance(m, MultiModalMessage)
+                assert isinstance(m, (TextMessage, MultiModalMessage, ToolCallSummaryMessage))
                 context.append(UserMessage(content=m.content, source=m.source))
         return context
 

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

@@ -1,5 +1,9 @@
 from typing import Any, Callable, List, Mapping
 
+from autogen_core import Component, ComponentModel
+from pydantic import BaseModel
+from typing_extensions import Self
+
 from ...base import ChatAgent, TerminationCondition
 from ...messages import AgentEvent, ChatMessage
 from ...state import RoundRobinManagerState
@@ -61,7 +65,15 @@ async def select_speaker(self, thread: List[AgentEvent | ChatMessage]) -> str:
         return current_speaker
 
 
-class RoundRobinGroupChat(BaseGroupChat):
+class RoundRobinGroupChatConfig(BaseModel):
+    """The declarative configuration RoundRobinGroupChat."""
+
+    participants: List[ComponentModel]
+    termination_condition: ComponentModel | None = None
+    max_turns: int | None = None
+
+
+class RoundRobinGroupChat(BaseGroupChat, Component[RoundRobinGroupChatConfig]):
     """A team that runs a group chat with participants taking turns in a round-robin fashion
     to publish a message to all.
 
@@ -133,6 +145,9 @@ async def main() -> None:
             asyncio.run(main())
     """
 
+    component_config_schema = RoundRobinGroupChatConfig
+    component_provider_override = "autogen_agentchat.teams.RoundRobinGroupChat"
+
     def __init__(
         self,
         participants: List[ChatAgent],
@@ -166,3 +181,20 @@ def _factory() -> RoundRobinGroupChatManager:
             )
 
         return _factory
+
+    def _to_config(self) -> RoundRobinGroupChatConfig:
+        participants = [participant.dump_component() for participant in self._participants]
+        termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
+        return RoundRobinGroupChatConfig(
+            participants=participants,
+            termination_condition=termination_condition,
+            max_turns=self._max_turns,
+        )
+
+    @classmethod
+    def _from_config(cls, config: RoundRobinGroupChatConfig) -> Self:
+        participants = [ChatAgent.load_component(participant) for participant in config.participants]
+        termination_condition = (
+            TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
+        )
+        return cls(participants, termination_condition=termination_condition, max_turns=config.max_turns)