add swarm manager option

ag2ai · Dec 16, 2024 · e0eaa15 · e0eaa15
1 parent 190a93d
commit e0eaa15
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/autogen/agentchat/contrib/swarm_agent.py b/autogen/agentchat/contrib/swarm_agent.py
@@ -32,6 +32,7 @@ class AfterWorkOption(Enum):
     TERMINATE = "TERMINATE"
     REVERT_TO_USER = "REVERT_TO_USER"
     STAY = "STAY"
+    SWARM_MANAGER = "SWARM_MANAGER"
 
 
 @dataclass
@@ -209,6 +210,8 @@ def determine_next_agent(last_speaker: SwarmAgent, groupchat: GroupChat):
                 return user_agent
             elif tmp_after_work == AfterWorkOption.STAY:
                 return last_speaker
+            elif tmp_after_work == AfterWorkOption.SWARM_MANAGER:
+                return "auto"
         elif isinstance(tmp_after_work, Callable):
             return tmp_after_work(last_speaker, groupchat.messages, groupchat, context_variables)
         else: