42 lines
1.6 KiB
Python
42 lines
1.6 KiB
Python
# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from typing import TYPE_CHECKING, Optional, Union
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from ..agent import Agent
|
|
|
|
if TYPE_CHECKING:
|
|
# Avoid circular import
|
|
from ..groupchat import GroupChat
|
|
|
|
|
|
class SpeakerSelectionResult(BaseModel):
|
|
"""Represents a speaker selection result that will be returned to GroupChat._prepare_and_select_agents to determine the next speaker.
|
|
|
|
This class can return an Agent, a None to end the conversation, or a string for a speaker selection method.
|
|
"""
|
|
|
|
terminate: Optional[bool] = None
|
|
agent_name: Optional[str] = None
|
|
speaker_selection_method: Optional[str] = None
|
|
|
|
def get_speaker_selection_result(self, groupchat: "GroupChat") -> Optional[Union[Agent, str]]:
|
|
"""Get the speaker selection result. If None, the conversation will end."""
|
|
if self.agent_name is not None:
|
|
# Find the agent by name in the groupchat
|
|
for agent in groupchat.agents:
|
|
if agent.name == self.agent_name:
|
|
return agent
|
|
raise ValueError(f"Agent '{self.agent_name}' not found in groupchat.")
|
|
elif self.speaker_selection_method is not None:
|
|
return self.speaker_selection_method
|
|
elif self.terminate is not None and self.terminate:
|
|
return None
|
|
else:
|
|
raise ValueError(
|
|
"Unable to establish speaker selection result. No terminate, agent, or speaker selection method provided."
|
|
)
|