Agentchat move termination (#3992)

This commit is contained in:
Victor Dibia 2024-10-29 18:37:26 -07:00 committed by GitHub
parent 0f4dd0cc6d
commit 75b00e76e1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 40 additions and 28 deletions

View File

@ -48,7 +48,7 @@ class MaxMessageTermination(TerminationCondition):
self._message_count += len(messages) self._message_count += len(messages)
if self._message_count >= self._max_messages: if self._message_count >= self._max_messages:
return StopMessage( return StopMessage(
content=f"Maximal number of messages {self._max_messages} reached, current message count: {self._message_count}", content=f"Maximum number of messages {self._max_messages} reached, current message count: {self._message_count}",
source="MaxMessageTermination", source="MaxMessageTermination",
) )
return None return None

View File

@ -28,7 +28,12 @@ class BaseGroupChat(Team, ABC):
create a subclass of :class:`BaseGroupChat` that uses the group chat manager. create a subclass of :class:`BaseGroupChat` that uses the group chat manager.
""" """
def __init__(self, participants: List[ChatAgent], group_chat_manager_class: type[BaseGroupChatManager]): def __init__(
self,
participants: List[ChatAgent],
group_chat_manager_class: type[BaseGroupChatManager],
termination_condition: TerminationCondition | None = None,
):
if len(participants) == 0: if len(participants) == 0:
raise ValueError("At least one participant is required.") raise ValueError("At least one participant is required.")
if len(participants) != len(set(participant.name for participant in participants)): if len(participants) != len(set(participant.name for participant in participants)):
@ -36,6 +41,7 @@ class BaseGroupChat(Team, ABC):
self._participants = participants self._participants = participants
self._team_id = str(uuid.uuid4()) self._team_id = str(uuid.uuid4())
self._base_group_chat_manager_class = group_chat_manager_class self._base_group_chat_manager_class = group_chat_manager_class
self._termination_condition = termination_condition
@abstractmethod @abstractmethod
def _create_group_chat_manager_factory( def _create_group_chat_manager_factory(
@ -109,7 +115,7 @@ class BaseGroupChat(Team, ABC):
group_topic_type=group_topic_type, group_topic_type=group_topic_type,
participant_topic_types=participant_topic_types, participant_topic_types=participant_topic_types,
participant_descriptions=participant_descriptions, participant_descriptions=participant_descriptions,
termination_condition=termination_condition, termination_condition=termination_condition or self._termination_condition,
), ),
) )
# Add subscriptions for the group chat manager. # Add subscriptions for the group chat manager.

View File

@ -82,8 +82,12 @@ class RoundRobinGroupChat(BaseGroupChat):
""" """
def __init__(self, participants: List[ChatAgent]): def __init__(self, participants: List[ChatAgent], termination_condition: TerminationCondition | None = None):
super().__init__(participants, group_chat_manager_class=RoundRobinGroupChatManager) super().__init__(
participants,
termination_condition=termination_condition,
group_chat_manager_class=RoundRobinGroupChatManager,
)
def _create_group_chat_manager_factory( def _create_group_chat_manager_factory(
self, self,

View File

@ -140,7 +140,8 @@ class SelectorGroupChatManager(BaseGroupChatManager):
+ re.escape(name.replace("_", r"\_")) + re.escape(name.replace("_", r"\_"))
+ r")(?=\W)" + r")(?=\W)"
) )
count = len(re.findall(regex, f" {message_content} ")) # Pad the message to help with matching # Pad the message to help with matching
count = len(re.findall(regex, f" {message_content} "))
if count > 0: if count > 0:
mentions[name] = count mentions[name] = count
return mentions return mentions
@ -184,6 +185,7 @@ class SelectorGroupChat(BaseGroupChat):
participants: List[ChatAgent], participants: List[ChatAgent],
model_client: ChatCompletionClient, model_client: ChatCompletionClient,
*, *,
termination_condition: TerminationCondition | None = None,
selector_prompt: str = """You are in a role play game. The following roles are available: selector_prompt: str = """You are in a role play game. The following roles are available:
{roles}. {roles}.
Read the following conversation. Then select the next role from {participants} to play. Only return the role. Read the following conversation. Then select the next role from {participants} to play. Only return the role.
@ -194,7 +196,9 @@ Read the above conversation. Then select the next role from {participants} to pl
""", """,
allow_repeated_speaker: bool = False, allow_repeated_speaker: bool = False,
): ):
super().__init__(participants, group_chat_manager_class=SelectorGroupChatManager) super().__init__(
participants, termination_condition=termination_condition, group_chat_manager_class=SelectorGroupChatManager
)
# Validate the participants. # Validate the participants.
if len(participants) < 2: if len(participants) < 2:
raise ValueError("At least two participants are required for SelectorGroupChat.") raise ValueError("At least two participants are required for SelectorGroupChat.")

View File

@ -82,8 +82,10 @@ class Swarm(BaseGroupChat):
await team.run("What is bob's birthday?", termination_condition=MaxMessageTermination(3)) await team.run("What is bob's birthday?", termination_condition=MaxMessageTermination(3))
""" """
def __init__(self, participants: List[ChatAgent]): def __init__(self, participants: List[ChatAgent], termination_condition: TerminationCondition | None = None):
super().__init__(participants, group_chat_manager_class=SwarmGroupChatManager) super().__init__(
participants, termination_condition=termination_condition, group_chat_manager_class=SwarmGroupChatManager
)
def _create_group_chat_manager_factory( def _create_group_chat_manager_factory(
self, self,

View File

@ -37,18 +37,18 @@
"text": [ "text": [
"\n", "\n",
"--------------------------------------------------------------------------- \n", "--------------------------------------------------------------------------- \n",
"\u001b[91m[2024-10-23T12:15:51.582079]:\u001b[0m\n", "\u001b[91m[2024-10-29T15:48:06.329810]:\u001b[0m\n",
"\n", "\n",
"What is the weather in New York?\n", "What is the weather in New York?\n",
"--------------------------------------------------------------------------- \n", "--------------------------------------------------------------------------- \n",
"\u001b[91m[2024-10-23T12:15:52.745820], writing_agent:\u001b[0m\n", "\u001b[91m[2024-10-29T15:48:08.085839], weather_agent:\u001b[0m\n",
"\n", "\n",
"The weather in New York is currently 73 degrees and sunny. TERMINATE\n", "The weather in New York is 73 degrees and sunny.\n",
"--------------------------------------------------------------------------- \n", "--------------------------------------------------------------------------- \n",
"\u001b[91m[2024-10-23T12:15:52.746210], Termination:\u001b[0m\n", "\u001b[91m[2024-10-29T15:48:08.086180], Termination:\u001b[0m\n",
"\n", "\n",
"Maximal number of messages 1 reached, current message count: 1\n", "Maximum number of messages 2 reached, current message count: 2\n",
" TaskResult(messages=[TextMessage(source='user', content='What is the weather in New York?'), StopMessage(source='writing_agent', content='The weather in New York is currently 73 degrees and sunny. TERMINATE')])\n" " TaskResult(messages=[TextMessage(source='user', content='What is the weather in New York?'), TextMessage(source='weather_agent', content='The weather in New York is 73 degrees and sunny.')])\n"
] ]
} }
], ],
@ -56,13 +56,13 @@
"import logging\n", "import logging\n",
"\n", "\n",
"from autogen_agentchat import EVENT_LOGGER_NAME\n", "from autogen_agentchat import EVENT_LOGGER_NAME\n",
"from autogen_agentchat.agents import ToolUseAssistantAgent\n", "from autogen_agentchat.agents import AssistantAgent\n",
"from autogen_agentchat.logging import ConsoleLogHandler\n", "from autogen_agentchat.logging import ConsoleLogHandler\n",
"from autogen_agentchat.task import MaxMessageTermination\n", "from autogen_agentchat.task import MaxMessageTermination\n",
"from autogen_agentchat.teams import RoundRobinGroupChat\n", "from autogen_agentchat.teams import RoundRobinGroupChat\n",
"from autogen_core.components.tools import FunctionTool\n",
"from autogen_ext.models import OpenAIChatCompletionClient\n", "from autogen_ext.models import OpenAIChatCompletionClient\n",
"\n", "\n",
"# set up logging. You can define your own logger\n",
"logger = logging.getLogger(EVENT_LOGGER_NAME)\n", "logger = logging.getLogger(EVENT_LOGGER_NAME)\n",
"logger.addHandler(ConsoleLogHandler())\n", "logger.addHandler(ConsoleLogHandler())\n",
"logger.setLevel(logging.INFO)\n", "logger.setLevel(logging.INFO)\n",
@ -73,22 +73,18 @@
" return f\"The weather in {city} is 73 degrees and Sunny.\"\n", " return f\"The weather in {city} is 73 degrees and Sunny.\"\n",
"\n", "\n",
"\n", "\n",
"# wrap the tool for use with the agent\n",
"get_weather_tool = FunctionTool(get_weather, description=\"Get the weather for a city\")\n",
"\n",
"# define an agent\n", "# define an agent\n",
"weather_agent = ToolUseAssistantAgent(\n", "weather_agent = AssistantAgent(\n",
" name=\"writing_agent\",\n", " name=\"weather_agent\",\n",
" model_client=OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\"),\n", " model_client=OpenAIChatCompletionClient(model=\"gpt-4o-2024-08-06\"),\n",
" registered_tools=[get_weather_tool],\n", " tools=[get_weather],\n",
")\n", ")\n",
"\n", "\n",
"# add the agent to a team\n", "# add the agent to a team\n",
"agent_team = RoundRobinGroupChat([weather_agent])\n", "agent_team = RoundRobinGroupChat([weather_agent], termination_condition=MaxMessageTermination(max_messages=2))\n",
"# Note: if running in a Python file directly you'll need to use asyncio.run(agent_team.run(...)) instead of await agent_team.run(...)\n", "# Note: if running in a Python file directly you'll need to use asyncio.run(agent_team.run(...)) instead of await agent_team.run(...)\n",
"result = await agent_team.run(\n", "result = await agent_team.run(\n",
" task=\"What is the weather in New York?\",\n", " task=\"What is the weather in New York?\",\n",
" termination_condition=MaxMessageTermination(max_messages=1),\n",
")\n", ")\n",
"print(\"\\n\", result)" "print(\"\\n\", result)"
] ]
@ -97,7 +93,7 @@
"cell_type": "markdown", "cell_type": "markdown",
"metadata": {}, "metadata": {},
"source": [ "source": [
"The code snippet above introduces two high level concepts in AgentChat: `Agent` and `Team`. An Agent helps us define what actions are taken when a message is received. Specifically, we use the `ToolUseAssistantAgent` preset - an agent that can be given a function that it can then use to address tasks. A Team helps us define the rules for how agents interact with each other. In the `RoundRobinGroupChat` team, agents receive messages in a sequential round-robin fashion. " "The code snippet above introduces two high level concepts in AgentChat: `Agent` and `Team`. An Agent helps us define what actions are taken when a message is received. Specifically, we use the `AssistantAgent` preset - an agent that can be given tools (functions) that it can then use to address tasks. A Team helps us define the rules for how agents interact with each other. In the `RoundRobinGroupChat` team, agents receive messages in a sequential round-robin fashion. "
] ]
}, },
{ {
@ -113,7 +109,7 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": ".venv", "display_name": "agnext",
"language": "python", "language": "python",
"name": "python3" "name": "python3"
}, },
@ -127,7 +123,7 @@
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython3",
"version": "3.12.6" "version": "3.11.9"
} }
}, },
"nbformat": 4, "nbformat": 4,