From 7a2d53e95e0ba18549fd05fcefd4941ce24f2bf3 Mon Sep 17 00:00:00 2001
From: Victor Dibia <victor.dibia@gmail.com>
Date: Wed, 15 Jan 2025 16:28:49 -0800
Subject: [PATCH 1/8] initial pass on making group chats declarative
---
.../teams/_group_chat/_base_group_chat.py | 6 ++-
.../_magentic_one/_magentic_one_group_chat.py | 47 +++++++++++++++++-
.../_group_chat/_round_robin_group_chat.py | 33 ++++++++++++-
.../teams/_group_chat/_selector_group_chat.py | 49 ++++++++++++++++++-
.../teams/_group_chat/_swarm_group_chat.py | 33 +++++++++++--
.../tests/test_group_chat.py | 15 ++++++
6 files changed, 173 insertions(+), 10 deletions(-)
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
index d27865e0a4ab..43af3c79b56b 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
@@ -14,8 +14,10 @@
MessageContext,
SingleThreadedAgentRuntime,
TypeSubscription,
+ ComponentBase
)
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],
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
index f09904fb3322..3d65a856666e 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
@@ -1,8 +1,10 @@
import logging
from typing import Callable, List
+from typing_extensions import Self
+from pydantic import BaseModel
from autogen_core.models import ChatCompletionClient
-
+from autogen_core import Component, ComponentModel
from .... import EVENT_LOGGER_NAME, TRACE_LOGGER_NAME
from ....base import ChatAgent, TerminationCondition
from .._base_group_chat import BaseGroupChat
@@ -13,7 +15,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
+ max_turns: int | 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 +86,9 @@ async def main() -> None:
}
"""
+ component_config_schema = MagenticOneGroupChatConfig
+ component_provider_override = "autogen_agentchat.teams.MagenticOneGroupChat"
+
def __init__(
self,
participants: List[ChatAgent],
@@ -117,3 +133,30 @@ def _create_group_chat_manager_factory(
self._final_answer_prompt,
termination_condition,
)
+
+
+ def _to_config(self) -> MagenticOneGroupChatConfig:
+ participants= [] # [p.dump_component() for p 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 = [BaseGroupChat.load_component(participant) for participant in config.participants]
+ model_client = config.model_client.load_component()
+ termination_condition = config.termination_condition.load_component() 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,
+ )
\ No newline at end of file
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
index d6901f04c988..28af7794851e 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
@@ -1,11 +1,14 @@
from typing import Any, Callable, List, Mapping
+from pydantic import BaseModel
+
from ...base import ChatAgent, TerminationCondition
from ...messages import AgentEvent, ChatMessage
from ...state import RoundRobinManagerState
from ._base_group_chat import BaseGroupChat
from ._base_group_chat_manager import BaseGroupChatManager
-
+from autogen_core import ComponentModel, Component
+from typing_extensions import Self
class RoundRobinGroupChatManager(BaseGroupChatManager):
"""A group chat manager that selects the next speaker in a round-robin fashion."""
@@ -61,7 +64,14 @@ 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
+ max_turns: int | 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 +143,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 +179,19 @@ def _factory() -> RoundRobinGroupChatManager:
)
return _factory
+
+ def _to_config(self) -> RoundRobinGroupChatConfig:
+ # participants = [participant.dump_component() for participant in self._participants]
+ 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 = [BaseGroupChat.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)
\ No newline at end of file
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
index dcc399ee2972..2463b0b3a97b 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
@@ -1,8 +1,10 @@
import logging
import re
from typing import Any, Callable, Dict, List, Mapping, Sequence
-
+from pydantic import BaseModel
+from typing_extensions import Self
from autogen_core.models import ChatCompletionClient, SystemMessage
+from autogen_core import Component, ComponentModel
from ... import TRACE_LOGGER_NAME
from ...base import ChatAgent, TerminationCondition
@@ -183,8 +185,18 @@ def _mentioned_agents(self, message_content: str, agent_names: List[str]) -> Dic
mentions[name] = count
return mentions
+class SelectorGroupChatConfig(BaseModel):
+ """The declarative configuration for SelectorGroupChat."""
+
+ participants: List[ComponentModel]
+ model_client: ComponentModel
+ termination_condition: ComponentModel | None
+ max_turns: int | None
+ selector_prompt: str
+ allow_repeated_speaker: bool
+ # selector_func: ComponentModel | None
-class SelectorGroupChat(BaseGroupChat):
+class SelectorGroupChat(BaseGroupChat, Component[SelectorGroupChatConfig]):
"""A group chat team that have participants takes turn to publish a message
to all, using a ChatCompletion model to select the next speaker after each message.
@@ -321,6 +333,9 @@ def selector_func(messages: Sequence[AgentEvent | ChatMessage]) -> str | None:
asyncio.run(main())
"""
+ component_config_schema = SelectorGroupChatConfig
+ component_provider_override = "autogen_agentchat.teams.SelectorGroupChat"
+
def __init__(
self,
participants: List[ChatAgent],
@@ -381,3 +396,33 @@ def _create_group_chat_manager_factory(
self._allow_repeated_speaker,
self._selector_func,
)
+
+
+ def _to_config(self) -> SelectorGroupChatConfig:
+ return SelectorGroupChatConfig(
+ participants= [] , # [p.dump_component() for p in self.participants],
+ model_client=self._model_client.dump_component(),
+ termination_condition=self.termination_condition.dump_component() if self.termination_condition else None,
+ max_turns=self.max_turns,
+ selector_prompt=self._selector_prompt,
+ allow_repeated_speaker=self._allow_repeated_speaker,
+ # selector_func=self._selector_func.dump_component() if self._selector_func else None,
+ )
+
+ @classmethod
+
+ def _from_config(cls, config: SelectorGroupChatConfig) -> Self:
+ return cls(
+ participants = [BaseGroupChat.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,
+ max_turns=config.max_turns,
+ selector_prompt=config.selector_prompt,
+ allow_repeated_speaker=config.allow_repeated_speaker,
+ # selector_func=ComponentLoader.load_component(config.selector_func, Callable[[Sequence[AgentEvent | ChatMessage]], str | None])
+ # if config.selector_func
+ # else None,
+ )
+
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
index a31a693e086f..238faf7cd0e7 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
@@ -1,11 +1,13 @@
from typing import Any, Callable, List, Mapping
+from pydantic import BaseModel
+
from ...base import ChatAgent, TerminationCondition
from ...messages import AgentEvent, ChatMessage, HandoffMessage
from ...state import SwarmManagerState
from ._base_group_chat import BaseGroupChat
from ._base_group_chat_manager import BaseGroupChatManager
-
+from autogen_core import ComponentModel, Component
class SwarmGroupChatManager(BaseGroupChatManager):
"""A group chat manager that selects the next speaker based on handoff message only."""
@@ -91,8 +93,14 @@ async def load_state(self, state: Mapping[str, Any]) -> None:
self._current_turn = swarm_state.current_turn
self._current_speaker = swarm_state.current_speaker
+class SwarmConfig(BaseModel):
+ """The declarative configuration for Swarm."""
+
+ participants: List[ComponentModel]
+ termination_condition: ComponentModel | None
+ max_turns: int | None
-class Swarm(BaseGroupChat):
+class Swarm(BaseGroupChat, Component[SwarmConfig]):
"""A group chat team that selects the next speaker based on handoff message only.
The first participant in the list of participants is the initial speaker.
@@ -180,6 +188,9 @@ async def main() -> None:
asyncio.run(main())
"""
+ component_config_schema = SwarmConfig
+ component_provider_override = "autogen_agentchat.teams.Swarm"
+
def __init__(
self,
participants: List[ChatAgent],
@@ -216,4 +227,20 @@ def _factory() -> SwarmGroupChatManager:
max_turns,
)
- return _factory
+ return _factory
+
+
+ def _to_config(self) -> SwarmConfig:
+ participants = [] # [participant.dump_component() for participant in self._participants]
+ termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
+ return SwarmConfig(
+ participants=participants,
+ termination_condition=termination_condition,
+ max_turns=self._max_turns,
+ )
+
+ @classmethod
+ def _from_config(cls, config: SwarmConfig) -> "Swarm":
+ participants = [BaseGroupChat.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)
diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py
index 25194d6049af..0813e7ccb8b0 100644
--- a/python/packages/autogen-agentchat/tests/test_group_chat.py
+++ b/python/packages/autogen-agentchat/tests/test_group_chat.py
@@ -1082,3 +1082,18 @@ async def test_round_robin_group_chat_with_message_list() -> None:
# Test with empty message list
with pytest.raises(ValueError, match="Task list cannot be empty"):
await team.run(task=[])
+
+
+@pytest.mark.asyncio
+async def test_declarative_groupchats() -> None:
+ # Create a simple team with echo agents
+ agent1 = _EchoAgent("Agent1", "First agent")
+ agent2 = _EchoAgent("Agent2", "Second agent")
+ termination = MaxMessageTermination(4) # Stop after 4 messages
+ round_robin_team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination)
+
+ round_robin_team_config = round_robin_team.dump_component()
+ assert round_robin_team_config.provider == "autogen_agentchat.teams.RoundRobinGroupChat"
+
+ round_robin_team_loaded = RoundRobinGroupChat.load_component(round_robin_team_config)
+ assert round_robin_team_loaded.component_type == "team"
\ No newline at end of file
From 9d753afcb30205de813996e9b44872d1ee07e5e3 Mon Sep 17 00:00:00 2001
From: Victor Dibia <victor.dibia@gmail.com>
Date: Wed, 15 Jan 2025 21:33:25 -0800
Subject: [PATCH 2/8] update group chat tests
---
.../tests/test_group_chat.py | 68 ++++++++++++++++---
1 file changed, 60 insertions(+), 8 deletions(-)
diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py
index 0813e7ccb8b0..66eec1ad7f92 100644
--- a/python/packages/autogen-agentchat/tests/test_group_chat.py
+++ b/python/packages/autogen-agentchat/tests/test_group_chat.py
@@ -28,6 +28,7 @@
RoundRobinGroupChat,
SelectorGroupChat,
Swarm,
+ MagenticOneGroupChat
)
from autogen_agentchat.teams._group_chat._round_robin_group_chat import RoundRobinGroupChatManager
from autogen_agentchat.teams._group_chat._selector_group_chat import SelectorGroupChatManager
@@ -1085,15 +1086,66 @@ async def test_round_robin_group_chat_with_message_list() -> None:
@pytest.mark.asyncio
-async def test_declarative_groupchats() -> None:
- # Create a simple team with echo agents
+async def test_declarative_groupchats_with_config() -> None:
+ # Create basic agents and components for testing
agent1 = _EchoAgent("Agent1", "First agent")
agent2 = _EchoAgent("Agent2", "Second agent")
- termination = MaxMessageTermination(4) # Stop after 4 messages
- round_robin_team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination)
+ termination = MaxMessageTermination(4)
+ model_client = ReplayChatCompletionClient(["test"])
- round_robin_team_config = round_robin_team.dump_component()
- assert round_robin_team_config.provider == "autogen_agentchat.teams.RoundRobinGroupChat"
+ # Test round robin - verify config is preserved
+ round_robin = RoundRobinGroupChat(
+ participants=[agent1, agent2],
+ termination_condition=termination,
+ max_turns=5
+ )
+ config = round_robin.dump_component()
+ loaded = RoundRobinGroupChat.load_component(config)
+ assert loaded.dump_component() == config
- round_robin_team_loaded = RoundRobinGroupChat.load_component(round_robin_team_config)
- assert round_robin_team_loaded.component_type == "team"
\ No newline at end of file
+ # Test selector group chat - verify config is preserved
+ selector_prompt = "Custom selector prompt with {roles}, {participants}, {history}"
+ selector = SelectorGroupChat(
+ participants=[agent1, agent2],
+ model_client=model_client,
+ termination_condition=termination,
+ max_turns=10,
+ selector_prompt=selector_prompt,
+ allow_repeated_speaker=True
+ )
+ config = selector.dump_component()
+ loaded = SelectorGroupChat.load_component(config)
+ assert loaded.dump_component() == config
+
+ # Test swarm with handoff termination
+ handoff_termination = HandoffTermination(target="Agent2")
+ swarm = Swarm(
+ participants=[agent1, agent2],
+ termination_condition=handoff_termination,
+ max_turns=5
+ )
+ config = swarm.dump_component()
+ loaded = Swarm.load_component(config)
+ assert loaded.dump_component() == config
+
+ # Test MagenticOne with custom parameters
+ magentic = MagenticOneGroupChat(
+ participants=[agent1],
+ model_client=model_client,
+ max_turns=15,
+ max_stalls=5,
+ final_answer_prompt="Custom prompt"
+ )
+ config = magentic.dump_component()
+ loaded = MagenticOneGroupChat.load_component(config)
+ assert loaded.dump_component() == config
+
+ # Verify component types are correctly set for each
+ for team in [loaded, selector, swarm, magentic]:
+ assert team.component_type == "team"
+
+ # Verify provider strings are correctly set
+ assert round_robin.dump_component().provider == "autogen_agentchat.teams.RoundRobinGroupChat"
+ assert selector.dump_component().provider == "autogen_agentchat.teams.SelectorGroupChat"
+ assert swarm.dump_component().provider == "autogen_agentchat.teams.Swarm"
+ assert magentic.dump_component().provider == "autogen_agentchat.teams.MagenticOneGroupChat"
\ No newline at end of file
From c9182647625d5270f6a5ef83ca7f56c5fb563f66 Mon Sep 17 00:00:00 2001
From: Victor Dibia <victor.dibia@gmail.com>
Date: Sat, 18 Jan 2025 08:09:51 -0800
Subject: [PATCH 3/8] update impl to include participant serialization for all
teams
---
.../_magentic_one/_magentic_one_group_chat.py | 13 +++++++------
.../teams/_group_chat/_round_robin_group_chat.py | 11 ++++++-----
.../teams/_group_chat/_selector_group_chat.py | 13 +++++++------
.../teams/_group_chat/_swarm_group_chat.py | 9 +++++----
.../autogen-agentchat/tests/test_group_chat.py | 11 ++++++++---
5 files changed, 33 insertions(+), 24 deletions(-)
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
index 3d65a856666e..25262bafd435 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
@@ -5,6 +5,7 @@
from autogen_core.models import ChatCompletionClient
from autogen_core import Component, ComponentModel
+from ....agents import BaseChatAgent
from .... import EVENT_LOGGER_NAME, TRACE_LOGGER_NAME
from ....base import ChatAgent, TerminationCondition
from .._base_group_chat import BaseGroupChat
@@ -21,8 +22,8 @@ class MagenticOneGroupChatConfig(BaseModel):
participants: List[ComponentModel]
model_client: ComponentModel
- termination_condition: ComponentModel | None
- max_turns: int | None
+ termination_condition: ComponentModel | None = None
+ max_turns: int | None = None
max_stalls: int
final_answer_prompt: str
@@ -136,7 +137,7 @@ def _create_group_chat_manager_factory(
def _to_config(self) -> MagenticOneGroupChatConfig:
- participants= [] # [p.dump_component() for p in self.participants],
+ 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,
@@ -149,9 +150,9 @@ def _to_config(self) -> MagenticOneGroupChatConfig:
@classmethod
def _from_config(cls, config: MagenticOneGroupChatConfig) -> Self:
- participants = [BaseGroupChat.load_component(participant) for participant in config.participants]
- model_client = config.model_client.load_component()
- termination_condition = config.termination_condition.load_component() if config.termination_condition else None
+ participants = [BaseChatAgent.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,
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
index 28af7794851e..94ee22ed2966 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
@@ -10,6 +10,8 @@
from autogen_core import ComponentModel, Component
from typing_extensions import Self
+from ...agents import BaseChatAgent
+
class RoundRobinGroupChatManager(BaseGroupChatManager):
"""A group chat manager that selects the next speaker in a round-robin fashion."""
@@ -68,8 +70,8 @@ class RoundRobinGroupChatConfig(BaseModel):
"""The declarative configuration RoundRobinGroupChat."""
participants: List[ComponentModel]
- termination_condition: ComponentModel | None
- max_turns: int | None
+ 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
@@ -181,8 +183,7 @@ def _factory() -> RoundRobinGroupChatManager:
return _factory
def _to_config(self) -> RoundRobinGroupChatConfig:
- # participants = [participant.dump_component() for participant in self._participants]
- participants = []
+ 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,
@@ -192,6 +193,6 @@ def _to_config(self) -> RoundRobinGroupChatConfig:
@classmethod
def _from_config(cls, config: RoundRobinGroupChatConfig) -> Self:
- participants = [BaseGroupChat.load_component(participant) for participant in config.participants]
+ participants = [BaseChatAgent.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)
\ No newline at end of file
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
index 2463b0b3a97b..16657d96e336 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
@@ -7,6 +7,7 @@
from autogen_core import Component, ComponentModel
from ... import TRACE_LOGGER_NAME
+from ...agents import BaseChatAgent
from ...base import ChatAgent, TerminationCondition
from ...messages import (
AgentEvent,
@@ -190,8 +191,8 @@ class SelectorGroupChatConfig(BaseModel):
participants: List[ComponentModel]
model_client: ComponentModel
- termination_condition: ComponentModel | None
- max_turns: int | None
+ termination_condition: ComponentModel | None = None
+ max_turns: int | None = None
selector_prompt: str
allow_repeated_speaker: bool
# selector_func: ComponentModel | None
@@ -400,10 +401,10 @@ def _create_group_chat_manager_factory(
def _to_config(self) -> SelectorGroupChatConfig:
return SelectorGroupChatConfig(
- participants= [] , # [p.dump_component() for p in self.participants],
+ participants = [participant.dump_component() for participant in self._participants] ,
model_client=self._model_client.dump_component(),
- termination_condition=self.termination_condition.dump_component() if self.termination_condition else None,
- max_turns=self.max_turns,
+ termination_condition=self._termination_condition.dump_component() if self._termination_condition else None,
+ max_turns=self._max_turns,
selector_prompt=self._selector_prompt,
allow_repeated_speaker=self._allow_repeated_speaker,
# selector_func=self._selector_func.dump_component() if self._selector_func else None,
@@ -413,7 +414,7 @@ def _to_config(self) -> SelectorGroupChatConfig:
def _from_config(cls, config: SelectorGroupChatConfig) -> Self:
return cls(
- participants = [BaseGroupChat.load_component(participant) for participant in config.participants],
+ participants = [BaseChatAgent.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
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
index 238faf7cd0e7..b8725a81e0f0 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
@@ -2,6 +2,7 @@
from pydantic import BaseModel
+from ...agents import BaseChatAgent
from ...base import ChatAgent, TerminationCondition
from ...messages import AgentEvent, ChatMessage, HandoffMessage
from ...state import SwarmManagerState
@@ -97,8 +98,8 @@ class SwarmConfig(BaseModel):
"""The declarative configuration for Swarm."""
participants: List[ComponentModel]
- termination_condition: ComponentModel | None
- max_turns: int | None
+ termination_condition: ComponentModel | None = None
+ max_turns: int | None = None
class Swarm(BaseGroupChat, Component[SwarmConfig]):
"""A group chat team that selects the next speaker based on handoff message only.
@@ -231,7 +232,7 @@ def _factory() -> SwarmGroupChatManager:
def _to_config(self) -> SwarmConfig:
- participants = [] # [participant.dump_component() for participant in self._participants]
+ participants = [participant.dump_component() for participant in self._participants]
termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
return SwarmConfig(
participants=participants,
@@ -241,6 +242,6 @@ def _to_config(self) -> SwarmConfig:
@classmethod
def _from_config(cls, config: SwarmConfig) -> "Swarm":
- participants = [BaseGroupChat.load_component(participant) for participant in config.participants]
+ participants = [BaseChatAgent.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)
diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py
index 66eec1ad7f92..483bc75e83b7 100644
--- a/python/packages/autogen-agentchat/tests/test_group_chat.py
+++ b/python/packages/autogen-agentchat/tests/test_group_chat.py
@@ -1088,10 +1088,15 @@ async def test_round_robin_group_chat_with_message_list() -> None:
@pytest.mark.asyncio
async def test_declarative_groupchats_with_config() -> None:
# Create basic agents and components for testing
- agent1 = _EchoAgent("Agent1", "First agent")
- agent2 = _EchoAgent("Agent2", "Second agent")
+ agent1 = AssistantAgent(
+ "agent_1", model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key=""),
+ handoffs=["agent_2"],
+ )
+ agent2 = AssistantAgent(
+ "agent_2", model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key="")
+ )
termination = MaxMessageTermination(4)
- model_client = ReplayChatCompletionClient(["test"])
+ model_client = OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key="")
# Test round robin - verify config is preserved
round_robin = RoundRobinGroupChat(
From 451757f959d69aba99e7f9bc4725415dd62eb4e4 Mon Sep 17 00:00:00 2001
From: Victor Dibia <victor.dibia@gmail.com>
Date: Sat, 18 Jan 2025 20:25:29 -0800
Subject: [PATCH 4/8] v1 making soc declarative
---
.../agents/_society_of_mind_agent.py | 41 +++++++++++++++++--
1 file changed, 38 insertions(+), 3 deletions(-)
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
index 713b518bd7df..a6c5d6a4afbe 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
@@ -1,11 +1,12 @@
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 autogen_agentchat.base import Response
from autogen_agentchat.state import SocietyOfMindAgentState
-
+from typing_extensions import Self
+from pydantic import BaseModel
from ..base import TaskResult, Team
from ..messages import (
AgentEvent,
@@ -16,7 +17,16 @@
from ._base_chat_agent import BaseChatAgent
-class SocietyOfMindAgent(BaseChatAgent):
+class SocietyOfMindAgentConfig(BaseModel):
+ """The declarative configuration for a SocietyOfMindAgent."""
+
+ 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 +84,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 +186,25 @@ 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(
+ 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:
+ team = Team.load_component(config.team)
+ model_client = ChatCompletionClient.load_component(config.model_client)
+ return cls(
+ team=team,
+ model_client=model_client,
+ description=config.description,
+ instruction=config.instruction,
+ response_prompt=config.response_prompt,
+ )
\ No newline at end of file
From 2b9196ee185316b5667f365ee8a0a0028cf2de4d Mon Sep 17 00:00:00 2001
From: Victor Dibia <victor.dibia@gmail.com>
Date: Wed, 22 Jan 2025 14:16:00 -0800
Subject: [PATCH 5/8] update memory test
---
python/packages/autogen-core/tests/test_memory.py | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/python/packages/autogen-core/tests/test_memory.py b/python/packages/autogen-core/tests/test_memory.py
index ebe1866ed15d..5948cc014bc4 100644
--- a/python/packages/autogen-core/tests/test_memory.py
+++ b/python/packages/autogen-core/tests/test_memory.py
@@ -21,10 +21,10 @@ def test_memory_protocol_attributes() -> None:
assert hasattr(Memory, "close")
-def test_memory_protocol_runtime_checkable() -> None:
- """Test that Memory protocol is properly runtime-checkable."""
+def test_memory_abc_implementation() -> None:
+ """Test that Memory ABC is properly implemented."""
- class ValidMemory:
+ class ValidMemory(Memory):
@property
def name(self) -> str:
return "test"
From a61c3cff93fa748f7124b7bdf8374c0f914cd2da Mon Sep 17 00:00:00 2001
From: Victor Dibia <victor.dibia@gmail.com>
Date: Wed, 22 Jan 2025 16:15:09 -0800
Subject: [PATCH 6/8] update chatagent and team base classes
---
.../agents/_society_of_mind_agent.py | 17 ++---
.../src/autogen_agentchat/base/_chat_agent.py | 7 ++-
.../src/autogen_agentchat/base/_team.py | 6 +-
.../teams/_group_chat/_base_group_chat.py | 2 +-
.../_magentic_one/_magentic_one_group_chat.py | 20 +++---
.../_group_chat/_round_robin_group_chat.py | 12 ++--
.../teams/_group_chat/_selector_group_chat.py | 11 ++--
.../teams/_group_chat/_swarm_group_chat.py | 17 ++---
.../tests/test_group_chat.py | 62 +++++++------------
.../tests/test_society_of_mind_agent.py | 10 +++
10 files changed, 87 insertions(+), 77 deletions(-)
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
index a6c5d6a4afbe..8f1a950e07a5 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
@@ -1,6 +1,6 @@
from typing import Any, AsyncGenerator, List, Mapping, Sequence
-from autogen_core import CancellationToken,Component, ComponentModel
+from autogen_core import CancellationToken, Component, ComponentModel
from autogen_core.models import ChatCompletionClient, LLMMessage, SystemMessage, UserMessage
from autogen_agentchat.base import Response
@@ -20,11 +20,13 @@
class SocietyOfMindAgentConfig(BaseModel):
"""The declarative configuration for a SocietyOfMindAgent."""
+ name: str
team: ComponentModel
model_client: ComponentModel
- description: str
- instruction: str
- response_prompt: str
+ description: str
+ instruction: str
+ response_prompt: str
+
class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
"""An agent that uses an inner team of agents to generate responses.
@@ -187,9 +189,9 @@ 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,
@@ -199,12 +201,13 @@ def _to_config(self) -> SocietyOfMindAgentConfig:
@classmethod
def _from_config(cls, config: SocietyOfMindAgentConfig) -> Self:
- team = Team.load_component(config.team)
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,
- )
\ No newline at end of file
+ )
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_chat_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_chat_agent.py
index 36a80efe019c..34f0f37ee6e2 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_chat_agent.py
+++ b/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:
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py
index 565ad225b866..d790130ae93d 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py
@@ -1,9 +1,13 @@
from typing import Any, Mapping
from abc import ABC, abstractmethod
from ._task import TaskRunner
+from pydantic import BaseModel
+from autogen_core import ComponentBase
-class Team(ABC, TaskRunner):
+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."""
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
index 43af3c79b56b..446c74ca14de 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
@@ -14,7 +14,7 @@
MessageContext,
SingleThreadedAgentRuntime,
TypeSubscription,
- ComponentBase
+ ComponentBase,
)
from autogen_core._closure_agent import ClosureContext
from pydantic import BaseModel
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
index 25262bafd435..7aa44023fb94 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
@@ -1,11 +1,10 @@
import logging
from typing import Callable, List
from typing_extensions import Self
-from pydantic import BaseModel
+from pydantic import BaseModel
from autogen_core.models import ChatCompletionClient
from autogen_core import Component, ComponentModel
-from ....agents import BaseChatAgent
from .... import EVENT_LOGGER_NAME, TRACE_LOGGER_NAME
from ....base import ChatAgent, TerminationCondition
from .._base_group_chat import BaseGroupChat
@@ -16,17 +15,17 @@
event_logger = logging.getLogger(EVENT_LOGGER_NAME)
-
class MagenticOneGroupChatConfig(BaseModel):
"""The declarative configuration for a MagenticOneGroupChat."""
-
+
participants: List[ComponentModel]
- model_client: 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.
@@ -135,9 +134,8 @@ def _create_group_chat_manager_factory(
termination_condition,
)
-
def _to_config(self) -> MagenticOneGroupChatConfig:
- participants= [participant.dump_component() for participant in self._participants]
+ 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,
@@ -150,9 +148,11 @@ def _to_config(self) -> MagenticOneGroupChatConfig:
@classmethod
def _from_config(cls, config: MagenticOneGroupChatConfig) -> Self:
- participants = [BaseChatAgent.load_component(participant) for participant in config.participants]
+ 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
+ termination_condition = (
+ TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
+ )
return cls(
participants,
model_client,
@@ -160,4 +160,4 @@ def _from_config(cls, config: MagenticOneGroupChatConfig) -> Self:
max_turns=config.max_turns,
max_stalls=config.max_stalls,
final_answer_prompt=config.final_answer_prompt,
- )
\ No newline at end of file
+ )
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
index 94ee22ed2966..83b6c5909381 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
@@ -10,7 +10,6 @@
from autogen_core import ComponentModel, Component
from typing_extensions import Self
-from ...agents import BaseChatAgent
class RoundRobinGroupChatManager(BaseGroupChatManager):
"""A group chat manager that selects the next speaker in a round-robin fashion."""
@@ -73,6 +72,7 @@ class RoundRobinGroupChatConfig(BaseModel):
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.
@@ -183,7 +183,7 @@ def _factory() -> RoundRobinGroupChatManager:
return _factory
def _to_config(self) -> RoundRobinGroupChatConfig:
- participants = [participant.dump_component() for participant in self._participants]
+ 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,
@@ -193,6 +193,8 @@ def _to_config(self) -> RoundRobinGroupChatConfig:
@classmethod
def _from_config(cls, config: RoundRobinGroupChatConfig) -> Self:
- participants = [BaseChatAgent.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)
\ No newline at end of file
+ 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)
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
index 16657d96e336..0a3b7fd85eb5 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
@@ -186,6 +186,7 @@ def _mentioned_agents(self, message_content: str, agent_names: List[str]) -> Dic
mentions[name] = count
return mentions
+
class SelectorGroupChatConfig(BaseModel):
"""The declarative configuration for SelectorGroupChat."""
@@ -197,6 +198,7 @@ class SelectorGroupChatConfig(BaseModel):
allow_repeated_speaker: bool
# selector_func: ComponentModel | None
+
class SelectorGroupChat(BaseGroupChat, Component[SelectorGroupChatConfig]):
"""A group chat team that have participants takes turn to publish a message
to all, using a ChatCompletion model to select the next speaker after each message.
@@ -398,10 +400,9 @@ def _create_group_chat_manager_factory(
self._selector_func,
)
-
def _to_config(self) -> SelectorGroupChatConfig:
return SelectorGroupChatConfig(
- participants = [participant.dump_component() for participant in self._participants] ,
+ participants=[participant.dump_component() for participant in self._participants],
model_client=self._model_client.dump_component(),
termination_condition=self._termination_condition.dump_component() if self._termination_condition else None,
max_turns=self._max_turns,
@@ -411,10 +412,9 @@ def _to_config(self) -> SelectorGroupChatConfig:
)
@classmethod
-
- def _from_config(cls, config: SelectorGroupChatConfig) -> Self:
+ def _from_config(cls, config: SelectorGroupChatConfig) -> Self:
return cls(
- participants = [BaseChatAgent.load_component(participant) for participant in config.participants],
+ participants=[BaseChatAgent.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
@@ -426,4 +426,3 @@ def _from_config(cls, config: SelectorGroupChatConfig) -> Self:
# if config.selector_func
# else None,
)
-
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
index b8725a81e0f0..ec77cb8e3555 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
@@ -2,7 +2,6 @@
from pydantic import BaseModel
-from ...agents import BaseChatAgent
from ...base import ChatAgent, TerminationCondition
from ...messages import AgentEvent, ChatMessage, HandoffMessage
from ...state import SwarmManagerState
@@ -10,6 +9,7 @@
from ._base_group_chat_manager import BaseGroupChatManager
from autogen_core import ComponentModel, Component
+
class SwarmGroupChatManager(BaseGroupChatManager):
"""A group chat manager that selects the next speaker based on handoff message only."""
@@ -94,6 +94,7 @@ async def load_state(self, state: Mapping[str, Any]) -> None:
self._current_turn = swarm_state.current_turn
self._current_speaker = swarm_state.current_speaker
+
class SwarmConfig(BaseModel):
"""The declarative configuration for Swarm."""
@@ -101,6 +102,7 @@ class SwarmConfig(BaseModel):
termination_condition: ComponentModel | None = None
max_turns: int | None = None
+
class Swarm(BaseGroupChat, Component[SwarmConfig]):
"""A group chat team that selects the next speaker based on handoff message only.
@@ -228,11 +230,10 @@ def _factory() -> SwarmGroupChatManager:
max_turns,
)
- return _factory
-
+ return _factory
def _to_config(self) -> SwarmConfig:
- participants = [participant.dump_component() for participant in self._participants]
+ participants = [participant.dump_component() for participant in self._participants]
termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
return SwarmConfig(
participants=participants,
@@ -240,8 +241,10 @@ def _to_config(self) -> SwarmConfig:
max_turns=self._max_turns,
)
- @classmethod
+ @classmethod
def _from_config(cls, config: SwarmConfig) -> "Swarm":
- participants = [BaseChatAgent.load_component(participant) for participant in config.participants]
- termination_condition = TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
+ 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)
diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py
index 157cffbb24de..3a402b4d07fd 100644
--- a/python/packages/autogen-agentchat/tests/test_group_chat.py
+++ b/python/packages/autogen-agentchat/tests/test_group_chat.py
@@ -24,12 +24,7 @@
ToolCallRequestEvent,
ToolCallSummaryMessage,
)
-from autogen_agentchat.teams import (
- RoundRobinGroupChat,
- SelectorGroupChat,
- Swarm,
- MagenticOneGroupChat
-)
+from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat, Swarm, MagenticOneGroupChat
from autogen_agentchat.teams._group_chat._round_robin_group_chat import RoundRobinGroupChatManager
from autogen_agentchat.teams._group_chat._selector_group_chat import SelectorGroupChatManager
from autogen_agentchat.teams._group_chat._swarm_group_chat import SwarmGroupChatManager
@@ -1225,22 +1220,17 @@ async def test_round_robin_group_chat_with_message_list() -> None:
@pytest.mark.asyncio
async def test_declarative_groupchats_with_config() -> None:
# Create basic agents and components for testing
- agent1 = AssistantAgent(
- "agent_1", model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key=""),
- handoffs=["agent_2"],
- )
- agent2 = AssistantAgent(
- "agent_2", model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key="")
+ agent1 = AssistantAgent(
+ "agent_1",
+ model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key=""),
+ handoffs=["agent_2"],
)
+ agent2 = AssistantAgent("agent_2", model_client=OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key=""))
termination = MaxMessageTermination(4)
model_client = OpenAIChatCompletionClient(model="gpt-4o-2024-05-13", api_key="")
# Test round robin - verify config is preserved
- round_robin = RoundRobinGroupChat(
- participants=[agent1, agent2],
- termination_condition=termination,
- max_turns=5
- )
+ round_robin = RoundRobinGroupChat(participants=[agent1, agent2], termination_condition=termination, max_turns=5)
config = round_robin.dump_component()
loaded = RoundRobinGroupChat.load_component(config)
assert loaded.dump_component() == config
@@ -1253,41 +1243,37 @@ async def test_declarative_groupchats_with_config() -> None:
termination_condition=termination,
max_turns=10,
selector_prompt=selector_prompt,
- allow_repeated_speaker=True
+ allow_repeated_speaker=True,
)
- config = selector.dump_component()
- loaded = SelectorGroupChat.load_component(config)
- assert loaded.dump_component() == config
-
+ selector_config = selector.dump_component()
+ selector_loaded = SelectorGroupChat.load_component(selector_config)
+ assert selector_loaded.dump_component() == selector_config
+
# Test swarm with handoff termination
handoff_termination = HandoffTermination(target="Agent2")
- swarm = Swarm(
- participants=[agent1, agent2],
- termination_condition=handoff_termination,
- max_turns=5
- )
- config = swarm.dump_component()
- loaded = Swarm.load_component(config)
- assert loaded.dump_component() == config
-
+ swarm = Swarm(participants=[agent1, agent2], termination_condition=handoff_termination, max_turns=5)
+ swarm_config = swarm.dump_component()
+ swarm_loaded = Swarm.load_component(swarm_config)
+ assert swarm_loaded.dump_component() == swarm_config
+
# Test MagenticOne with custom parameters
magentic = MagenticOneGroupChat(
participants=[agent1],
model_client=model_client,
max_turns=15,
max_stalls=5,
- final_answer_prompt="Custom prompt"
+ final_answer_prompt="Custom prompt",
)
- config = magentic.dump_component()
- loaded = MagenticOneGroupChat.load_component(config)
- assert loaded.dump_component() == config
+ magentic_config = magentic.dump_component()
+ magentic_loaded = MagenticOneGroupChat.load_component(magentic_config)
+ assert magentic_loaded.dump_component() == magentic_config
# Verify component types are correctly set for each
for team in [loaded, selector, swarm, magentic]:
assert team.component_type == "team"
-
+
# Verify provider strings are correctly set
assert round_robin.dump_component().provider == "autogen_agentchat.teams.RoundRobinGroupChat"
assert selector.dump_component().provider == "autogen_agentchat.teams.SelectorGroupChat"
- assert swarm.dump_component().provider == "autogen_agentchat.teams.Swarm"
- assert magentic.dump_component().provider == "autogen_agentchat.teams.MagenticOneGroupChat"
\ No newline at end of file
+ assert swarm.dump_component().provider == "autogen_agentchat.teams.Swarm"
+ assert magentic.dump_component().provider == "autogen_agentchat.teams.MagenticOneGroupChat"
diff --git a/python/packages/autogen-agentchat/tests/test_society_of_mind_agent.py b/python/packages/autogen-agentchat/tests/test_society_of_mind_agent.py
index 9bf4713d9c43..f71ba67d5d7a 100644
--- a/python/packages/autogen-agentchat/tests/test_society_of_mind_agent.py
+++ b/python/packages/autogen-agentchat/tests/test_society_of_mind_agent.py
@@ -89,3 +89,13 @@ async def test_society_of_mind_agent(monkeypatch: pytest.MonkeyPatch) -> None:
await society_of_mind_agent2.load_state(state)
state2 = await society_of_mind_agent2.save_state()
assert state == state2
+
+ # Test serialization.
+
+ soc_agent_config = society_of_mind_agent.dump_component()
+ assert soc_agent_config.provider == "autogen_agentchat.agents.SocietyOfMindAgent"
+
+ # Test deserialization.
+ loaded_soc_agent = SocietyOfMindAgent.load_component(soc_agent_config)
+ assert isinstance(loaded_soc_agent, SocietyOfMindAgent)
+ assert loaded_soc_agent.name == "society_of_mind"
From c0ad84494614e70e9d7ac9f5a32ab480e5c89cc1 Mon Sep 17 00:00:00 2001
From: Victor Dibia <victor.dibia@gmail.com>
Date: Wed, 22 Jan 2025 16:21:38 -0800
Subject: [PATCH 7/8] update serialization doc notebook
---
.../serialize-components.ipynb | 56 ++++++++++++++++++-
1 file changed, 55 insertions(+), 1 deletion(-)
diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb
index 5a3855f48080..aa30b97d6eeb 100644
--- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb
+++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb
@@ -15,7 +15,7 @@
"\n",
"We will be implementing declarative support for the following components:\n",
"\n",
- "- Termination conditions ✔️\n",
+ "- Termination conditions\n",
"- Tools \n",
"- Agents \n",
"- Teams \n",
@@ -145,6 +145,60 @@
"\n",
"```"
]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Team Example\n",
+ "\n",
+ "In the example below, we will define a team in python, export this to a dictionary/json and also demonstrate how the team object can be loaded from the dictionary/json."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "{\"provider\":\"autogen_agentchat.teams.RoundRobinGroupChat\",\"component_type\":\"team\",\"version\":1,\"component_version\":1,\"description\":null,\"config\":{\"participants\":[{\"provider\":\"autogen_agentchat.agents.AssistantAgent\",\"component_type\":\"agent\",\"version\":1,\"component_version\":1,\"config\":{\"name\":\"assistant\",\"model_client\":{\"provider\":\"autogen_ext.models.openai.OpenAIChatCompletionClient\",\"component_type\":\"model\",\"version\":1,\"component_version\":1,\"config\":{\"model\":\"gpt-4o\"}},\"handoffs\":[{\"target\":\"flights_refunder\",\"description\":\"Handoff to flights_refunder.\",\"name\":\"transfer_to_flights_refunder\",\"message\":\"Transferred to flights_refunder, adopting the role of flights_refunder immediately.\"},{\"target\":\"user\",\"description\":\"Handoff to user.\",\"name\":\"transfer_to_user\",\"message\":\"Transferred to user, adopting the role of user immediately.\"}],\"model_context\":{\"provider\":\"autogen_core.model_context.UnboundedChatCompletionContext\",\"component_type\":\"chat_completion_context\",\"version\":1,\"component_version\":1,\"config\":{}},\"description\":\"An agent that provides assistance with ability to use tools.\",\"system_message\":\"Use tools to solve tasks.\",\"reflect_on_tool_use\":false,\"tool_call_summary_format\":\"{result}\"}}],\"termination_condition\":{\"provider\":\"autogen_agentchat.conditions.MaxMessageTermination\",\"component_type\":\"termination\",\"version\":1,\"component_version\":1,\"config\":{\"max_messages\":2}}}}\n"
+ ]
+ }
+ ],
+ "source": [
+ "from autogen_agentchat.agents import AssistantAgent, UserProxyAgent\n",
+ "from autogen_agentchat.teams import RoundRobinGroupChat\n",
+ "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
+ "from autogen_agentchat.conditions import MaxMessageTermination\n",
+ "\n",
+ "# Create an agent that uses the OpenAI GPT-4o model.\n",
+ "model_client = OpenAIChatCompletionClient(\n",
+ " model=\"gpt-4o\",\n",
+ " # api_key=\"YOUR_API_KEY\",\n",
+ ")\n",
+ "agent = AssistantAgent(\n",
+ " name=\"assistant\",\n",
+ " model_client=model_client,\n",
+ " handoffs=[\"flights_refunder\", \"user\"],\n",
+ " # tools=[], # serializing tools is not yet supported\n",
+ " system_message=\"Use tools to solve tasks.\",\n",
+ ")\n",
+ "\n",
+ "team = RoundRobinGroupChat(participants=[agent], termination_condition=MaxMessageTermination(2))\n",
+ "\n",
+ "team_config = team.dump_component() # dump component\n",
+ "print(team_config.model_dump_json())"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
}
],
"metadata": {
From 9764d1617ccf5d3d18e0f1bf3949afedf31f3dff Mon Sep 17 00:00:00 2001
From: Victor Dibia <victor.dibia@gmail.com>
Date: Wed, 22 Jan 2025 18:36:23 -0800
Subject: [PATCH 8/8] fomating updates
---
.../autogen_agentchat/agents/_society_of_mind_agent.py | 5 +++--
.../autogen-agentchat/src/autogen_agentchat/base/_team.py | 8 +++++---
.../teams/_group_chat/_base_group_chat.py | 2 +-
.../_group_chat/_magentic_one/_magentic_one_group_chat.py | 7 ++++---
.../teams/_group_chat/_round_robin_group_chat.py | 4 ++--
.../teams/_group_chat/_selector_group_chat.py | 5 +++--
.../teams/_group_chat/_swarm_group_chat.py | 2 +-
.../packages/autogen-agentchat/tests/test_group_chat.py | 2 +-
.../agentchat-user-guide/serialize-components.ipynb | 2 +-
.../autogen-core/src/autogen_core/memory/_base_memory.py | 2 +-
python/packages/autogen-core/tests/test_memory.py | 1 +
11 files changed, 23 insertions(+), 17 deletions(-)
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
index 8f1a950e07a5..c43c472c2915 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_society_of_mind_agent.py
@@ -2,11 +2,12 @@
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
-from typing_extensions import Self
-from pydantic import BaseModel
+
from ..base import TaskResult, Team
from ..messages import (
AgentEvent,
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py
index d790130ae93d..0d25edf6d26c 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/base/_team.py
@@ -1,8 +1,10 @@
-from typing import Any, Mapping
from abc import ABC, abstractmethod
-from ._task import TaskRunner
-from pydantic import BaseModel
+from typing import Any, Mapping
+
from autogen_core import ComponentBase
+from pydantic import BaseModel
+
+from ._task import TaskRunner
class Team(ABC, TaskRunner, ComponentBase[BaseModel]):
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
index 446c74ca14de..61e3783a80e5 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.py
@@ -11,10 +11,10 @@
AgentType,
CancellationToken,
ClosureAgent,
+ ComponentBase,
MessageContext,
SingleThreadedAgentRuntime,
TypeSubscription,
- ComponentBase,
)
from autogen_core._closure_agent import ClosureContext
from pydantic import BaseModel
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
index 7aa44023fb94..5ce7a71ebe21 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py
@@ -1,10 +1,11 @@
import logging
from typing import Callable, List
-from typing_extensions import Self
-from pydantic import BaseModel
-from autogen_core.models import ChatCompletionClient
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
from .._base_group_chat import BaseGroupChat
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
index 83b6c5909381..c7d34f0b3322 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_round_robin_group_chat.py
@@ -1,14 +1,14 @@
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
from ._base_group_chat import BaseGroupChat
from ._base_group_chat_manager import BaseGroupChatManager
-from autogen_core import ComponentModel, Component
-from typing_extensions import Self
class RoundRobinGroupChatManager(BaseGroupChatManager):
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
index 0a3b7fd85eb5..2cc3c02a645d 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_selector_group_chat.py
@@ -1,10 +1,11 @@
import logging
import re
from typing import Any, Callable, Dict, List, Mapping, Sequence
+
+from autogen_core import Component, ComponentModel
+from autogen_core.models import ChatCompletionClient, SystemMessage
from pydantic import BaseModel
from typing_extensions import Self
-from autogen_core.models import ChatCompletionClient, SystemMessage
-from autogen_core import Component, ComponentModel
from ... import TRACE_LOGGER_NAME
from ...agents import BaseChatAgent
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
index ec77cb8e3555..0ca02420537b 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_swarm_group_chat.py
@@ -1,5 +1,6 @@
from typing import Any, Callable, List, Mapping
+from autogen_core import Component, ComponentModel
from pydantic import BaseModel
from ...base import ChatAgent, TerminationCondition
@@ -7,7 +8,6 @@
from ...state import SwarmManagerState
from ._base_group_chat import BaseGroupChat
from ._base_group_chat_manager import BaseGroupChatManager
-from autogen_core import ComponentModel, Component
class SwarmGroupChatManager(BaseGroupChatManager):
diff --git a/python/packages/autogen-agentchat/tests/test_group_chat.py b/python/packages/autogen-agentchat/tests/test_group_chat.py
index 3a402b4d07fd..c04d7344029f 100644
--- a/python/packages/autogen-agentchat/tests/test_group_chat.py
+++ b/python/packages/autogen-agentchat/tests/test_group_chat.py
@@ -24,7 +24,7 @@
ToolCallRequestEvent,
ToolCallSummaryMessage,
)
-from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat, Swarm, MagenticOneGroupChat
+from autogen_agentchat.teams import MagenticOneGroupChat, RoundRobinGroupChat, SelectorGroupChat, Swarm
from autogen_agentchat.teams._group_chat._round_robin_group_chat import RoundRobinGroupChatManager
from autogen_agentchat.teams._group_chat._selector_group_chat import SelectorGroupChatManager
from autogen_agentchat.teams._group_chat._swarm_group_chat import SwarmGroupChatManager
diff --git a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb
index aa30b97d6eeb..16aecc399999 100644
--- a/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb
+++ b/python/packages/autogen-core/docs/src/user-guide/agentchat-user-guide/serialize-components.ipynb
@@ -170,9 +170,9 @@
],
"source": [
"from autogen_agentchat.agents import AssistantAgent, UserProxyAgent\n",
+ "from autogen_agentchat.conditions import MaxMessageTermination\n",
"from autogen_agentchat.teams import RoundRobinGroupChat\n",
"from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
- "from autogen_agentchat.conditions import MaxMessageTermination\n",
"\n",
"# Create an agent that uses the OpenAI GPT-4o model.\n",
"model_client = OpenAIChatCompletionClient(\n",
diff --git a/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py b/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py
index bfccafdf0a9e..2ae79b4106bb 100644
--- a/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py
+++ b/python/packages/autogen-core/src/autogen_core/memory/_base_memory.py
@@ -1,6 +1,6 @@
+from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, Dict, List, Union
-from abc import ABC, abstractmethod
from pydantic import BaseModel, ConfigDict
diff --git a/python/packages/autogen-core/tests/test_memory.py b/python/packages/autogen-core/tests/test_memory.py
index 7d0f6b6ea5b6..04054e1b2250 100644
--- a/python/packages/autogen-core/tests/test_memory.py
+++ b/python/packages/autogen-core/tests/test_memory.py
@@ -1,4 +1,5 @@
from typing import Any
+
import pytest
from autogen_core import CancellationToken
from autogen_core.memory import (