{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Selector Group Chat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` implements a team where participants take turns broadcasting messages to all other participants,\n",
    "with the next speaker selected by a generative model (e.g., an LLM) based on the shared context. \n",
    "This enables dynamic and context-aware multi-agent collaboration.\n",
    "\n",
    "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` provides several key features:\n",
    "- Model-based speaker selection\n",
    "- Configurable participant roles and descriptions\n",
    "- Optional prevention of consecutive turns by the same speaker\n",
    "- Customizable selection prompting\n",
    "- Customizable selection function to override the default model-based selection\n",
    "\n",
    "```{note}\n",
    "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` is a high-level API.\n",
    "If you need more control and customization that is not supported by this API,\n",
    "you can take a look at the [Group Chat Pattern](../../core-user-guide/design-patterns/group-chat.ipynb)\n",
    "in the Core API documentation and implement your own group chat logic.\n",
    "```\n",
    "\n",
    "## How does it work?\n",
    "\n",
    "{py:class}`~autogen_agentchat.teams.SelectorGroupChat` is a group chat similar to {py:class}`~autogen_agentchat.teams.RoundRobinGroupChat`,\n",
    "but with a model-based next speaker selection mechanism.\n",
    "When the team receives a task through {py:meth}`~autogen_agentchat.teams.BaseGroupChat.run` or {py:meth}`~autogen_agentchat.teams.BaseGroupChat.run_stream`,\n",
    "the following steps are executed:\n",
    "\n",
    "1. The team analyzes the current conversation context, including the conversation history and participants' {py:attr}`~autogen_agentchat.base.ChatAgent.name` and {py:attr}`~autogen_agentchat.base.ChatAgent.description` attributes, to determine the next speaker using a model. You can override the model by providing a custom selection function.\n",
    "2. The team prompts the selected speaker agent to provide a response, which is then **broadcasted** to all other participants.\n",
    "3. The termination condition is checked to determine if the conversation should end, if not, the process repeats from step 1.\n",
    "4. When the conversation ends, the team returns the {py:class}`~autogen_agentchat.base.TaskResult` containing the conversation history from this task.\n",
    "\n",
    "Once the team finishes the task, the conversation context is kept within the team and all participants, so the next task can continue from the previous conversation context.\n",
    "You can reset the conversation context by calling {py:meth}`~autogen_agentchat.teams.BaseGroupChat.reset`.\n",
    "\n",
    "In this section, we will demonstrate how to use {py:class}`~autogen_agentchat.teams.SelectorGroupChat` with a simple example for a web search and data analysis task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Web Search and Analysis Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Sequence\n",
    "\n",
    "from autogen_agentchat.agents import AssistantAgent\n",
    "from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination\n",
    "from autogen_agentchat.messages import AgentMessage\n",
    "from autogen_agentchat.teams import SelectorGroupChat\n",
    "from autogen_agentchat.ui import Console\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agents\n",
    "\n",
    "![Selector Group Chat](selector-group-chat.svg)\n",
    "\n",
    "This system uses three specialized agents:\n",
    "\n",
    "- **Planning Agent**: The strategic coordinator that breaks down complex tasks into manageable subtasks. \n",
    "- **Web Search Agent**: An information retrieval specialist that interfaces with the `search_web_tool`.\n",
    "- **Data Analyst Agent**: An agent specialist in performing calculations equipped with `percentage_change_tool`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The tools `search_web_tool` and `percentage_change_tool` are external tools that the agents can use to perform their tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: This example uses mock tools instead of real APIs for demonstration purposes\n",
    "def search_web_tool(query: str) -> str:\n",
    "    if \"2006-2007\" in query:\n",
    "        return \"\"\"Here are the total points scored by Miami Heat players in the 2006-2007 season:\n",
    "        Udonis Haslem: 844 points\n",
    "        Dwayne Wade: 1397 points\n",
    "        James Posey: 550 points\n",
    "        ...\n",
    "        \"\"\"\n",
    "    elif \"2007-2008\" in query:\n",
    "        return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\"\n",
    "    elif \"2008-2009\" in query:\n",
    "        return \"The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\"\n",
    "    return \"No data found.\"\n",
    "\n",
    "\n",
    "def percentage_change_tool(start: float, end: float) -> float:\n",
    "    return ((end - start) / start) * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create the specialized agents using the {py:class}`~autogen_agentchat.agents.AssistantAgent` class.\n",
    "It is important to note that the agents' {py:attr}`~autogen_agentchat.base.ChatAgent.name` and {py:attr}`~autogen_agentchat.base.ChatAgent.description` attributes are used by the model to determine the next speaker,\n",
    "so it is recommended to provide meaningful names and descriptions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_client = OpenAIChatCompletionClient(model=\"gpt-4o\")\n",
    "\n",
    "planning_agent = AssistantAgent(\n",
    "    \"PlanningAgent\",\n",
    "    description=\"An agent for planning tasks, this agent should be the first to engage when given a new task.\",\n",
    "    model_client=model_client,\n",
    "    system_message=\"\"\"\n",
    "    You are a planning agent.\n",
    "    Your job is to break down complex tasks into smaller, manageable subtasks.\n",
    "    Your team members are:\n",
    "        Web search agent: Searches for information\n",
    "        Data analyst: Performs calculations\n",
    "\n",
    "    You only plan and delegate tasks - you do not execute them yourself.\n",
    "\n",
    "    When assigning tasks, use this format:\n",
    "    1. <agent> : <task>\n",
    "\n",
    "    After all tasks are complete, summarize the findings and end with \"TERMINATE\".\n",
    "    \"\"\",\n",
    ")\n",
    "\n",
    "web_search_agent = AssistantAgent(\n",
    "    \"WebSearchAgent\",\n",
    "    description=\"A web search agent.\",\n",
    "    tools=[search_web_tool],\n",
    "    model_client=model_client,\n",
    "    system_message=\"\"\"\n",
    "    You are a web search agent.\n",
    "    Your only tool is search_tool - use it to find information.\n",
    "    You make only one search call at a time.\n",
    "    Once you have the results, you never do calculations based on them.\n",
    "    \"\"\",\n",
    ")\n",
    "\n",
    "data_analyst_agent = AssistantAgent(\n",
    "    \"DataAnalystAgent\",\n",
    "    description=\"A data analyst agent. Useful for performing calculations.\",\n",
    "    model_client=model_client,\n",
    "    tools=[percentage_change_tool],\n",
    "    system_message=\"\"\"\n",
    "    You are a data analyst.\n",
    "    Given the tasks you have been assigned, you should analyze the data and provide results using the tools provided.\n",
    "    \"\"\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Workflow\n",
    "\n",
    "1. The task is received by the {py:class}`~autogen_agentchat.teams.SelectorGroupChat` which, based on agent descriptions, selects the most appropriate agent to handle the initial task (typically the Planning Agent).\n",
    "\n",
    "2. The **Planning Agent** analyzes the task and breaks it down into subtasks, assigning each to the most appropriate agent using the format:\n",
    "   `<agent> : <task>`\n",
    "\n",
    "3. Based on the conversation context and agent descriptions, the {py:class}`~autogen_agent.teams.SelectorGroupChat` manager dynamically selects the next agent to handle their assigned subtask.\n",
    "\n",
    "4. The **Web Search Agent** performs searches one at a time, storing results in the shared conversation history.\n",
    "\n",
    "5. The **Data Analyst** processes the gathered information using available calculation tools when selected.\n",
    "\n",
    "6. The workflow continues with agents being dynamically selected until either:\n",
    "   - The Planning Agent determines all subtasks are complete and sends \"TERMINATE\"\n",
    "   - An alternative termination condition is met (e.g., a maximum number of messages)\n",
    "\n",
    "When defining your agents, make sure to include a helpful {py:attr}`~autogen_agentchat.base.ChatAgent.description` since this is used to decide which agent to select next."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create the team with two termination conditions:\n",
    "{py:class}`~autogen_agentchat.conditions.TextMentionTermination` to end the conversation when the Planning Agent sends \"TERMINATE\",\n",
    "and {py:class}`~autogen_agentchat.conditions.MaxMessageTermination` to limit the conversation to 25 messages to avoid infinite loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_mention_termination = TextMentionTermination(\"TERMINATE\")\n",
    "max_messages_termination = MaxMessageTermination(max_messages=25)\n",
    "termination = text_mention_termination | max_messages_termination\n",
    "\n",
    "team = SelectorGroupChat(\n",
    "    [planning_agent, web_search_agent, data_analyst_agent],\n",
    "    model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n",
    "    termination_condition=termination,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we run the team with a task to find information about an NBA player."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- user ----------\n",
      "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n",
      "---------- PlanningAgent ----------\n",
      "To address this request, we will divide the task into manageable subtasks. \n",
      "\n",
      "1. Web search agent: Identify the Miami Heat player with the highest points in the 2006-2007 season.\n",
      "2. Web search agent: Gather the total rebounds for the identified player during the 2007-2008 season.\n",
      "3. Web search agent: Gather the total rebounds for the identified player during the 2008-2009 season.\n",
      "4. Data analyst: Calculate the percentage change in total rebounds for the identified player between the 2007-2008 and 2008-2009 seasons.\n",
      "[Prompt tokens: 159, Completion tokens: 122]\n",
      "---------- WebSearchAgent ----------\n",
      "[FunctionCall(id='call_xdYlGP2lsqDeWdSiOlwOBNiO', arguments='{\"query\":\"Miami Heat highest points player 2006-2007 season\"}', name='search_web_tool')]\n",
      "[Prompt tokens: 271, Completion tokens: 26]\n",
      "---------- WebSearchAgent ----------\n",
      "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n        Udonis Haslem: 844 points\\n        Dwayne Wade: 1397 points\\n        James Posey: 550 points\\n        ...\\n        ', call_id='call_xdYlGP2lsqDeWdSiOlwOBNiO')]\n",
      "---------- WebSearchAgent ----------\n",
      "Tool calls:\n",
      "search_web_tool({\"query\":\"Miami Heat highest points player 2006-2007 season\"}) = Here are the total points scored by Miami Heat players in the 2006-2007 season:\n",
      "        Udonis Haslem: 844 points\n",
      "        Dwayne Wade: 1397 points\n",
      "        James Posey: 550 points\n",
      "        ...\n",
      "        \n",
      "---------- DataAnalystAgent ----------\n",
      "[FunctionCall(id='call_asjxKtAVGfqrYl0jMpvwCrvV', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='WebSearchAgent'), FunctionCall(id='call_8L91Kizt0KU6RNwUgvNx7S0s', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='WebSearchAgent')]\n",
      "[Prompt tokens: 345, Completion tokens: 68]\n",
      "---------- DataAnalystAgent ----------\n",
      "[FunctionExecutionResult(content=\"Error: The tool 'WebSearchAgent' is not available.\", call_id='call_asjxKtAVGfqrYl0jMpvwCrvV'), FunctionExecutionResult(content=\"Error: The tool 'WebSearchAgent' is not available.\", call_id='call_8L91Kizt0KU6RNwUgvNx7S0s')]\n",
      "---------- DataAnalystAgent ----------\n",
      "Tool calls:\n",
      "WebSearchAgent({\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}) = Error: The tool 'WebSearchAgent' is not available.\n",
      "WebSearchAgent({\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}) = Error: The tool 'WebSearchAgent' is not available.\n",
      "---------- WebSearchAgent ----------\n",
      "[FunctionCall(id='call_imvRJ2jhpPdovBbx8MFjlFVS', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_U30KVmFG1aeXPbqGJjDmJ6iJ', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')]\n",
      "[Prompt tokens: 445, Completion tokens: 70]\n",
      "---------- WebSearchAgent ----------\n",
      "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_imvRJ2jhpPdovBbx8MFjlFVS'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_U30KVmFG1aeXPbqGJjDmJ6iJ')]\n",
      "---------- WebSearchAgent ----------\n",
      "Tool calls:\n",
      "search_web_tool({\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}) = The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n",
      "search_web_tool({\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}) = The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n",
      "---------- DataAnalystAgent ----------\n",
      "[FunctionCall(id='call_CtAnvcbitN0JiwBfiLVzb5Do', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n",
      "[Prompt tokens: 562, Completion tokens: 20]\n",
      "---------- DataAnalystAgent ----------\n",
      "[FunctionExecutionResult(content='85.98130841121495', call_id='call_CtAnvcbitN0JiwBfiLVzb5Do')]\n",
      "---------- DataAnalystAgent ----------\n",
      "Tool calls:\n",
      "percentage_change_tool({\"start\":214,\"end\":398}) = 85.98130841121495\n",
      "---------- PlanningAgent ----------\n",
      "Summary of Findings:\n",
      "\n",
      "1. Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring a total of 1,397 points.\n",
      "2. Dwyane Wade's total rebounds during the 2007-2008 season were 214.\n",
      "3. Dwyane Wade's total rebounds during the 2008-2009 season were 398.\n",
      "4. The percentage change in Dwyane Wade's total rebounds between the 2007-2008 and 2008-2009 seasons was approximately 85.98%.\n",
      "\n",
      "TERMINATE\n",
      "[Prompt tokens: 590, Completion tokens: 122]\n",
      "---------- Summary ----------\n",
      "Number of messages: 15\n",
      "Finish reason: Text 'TERMINATE' mentioned\n",
      "Total prompt tokens: 2372\n",
      "Total completion tokens: 428\n",
      "Duration: 9.21 seconds\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=159, completion_tokens=122), content='To address this request, we will divide the task into manageable subtasks. \\n\\n1. Web search agent: Identify the Miami Heat player with the highest points in the 2006-2007 season.\\n2. Web search agent: Gather the total rebounds for the identified player during the 2007-2008 season.\\n3. Web search agent: Gather the total rebounds for the identified player during the 2008-2009 season.\\n4. Data analyst: Calculate the percentage change in total rebounds for the identified player between the 2007-2008 and 2008-2009 seasons.', type='TextMessage'), ToolCallMessage(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=271, completion_tokens=26), content=[FunctionCall(id='call_xdYlGP2lsqDeWdSiOlwOBNiO', arguments='{\"query\":\"Miami Heat highest points player 2006-2007 season\"}', name='search_web_tool')], type='ToolCallMessage'), ToolCallResultMessage(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n        Udonis Haslem: 844 points\\n        Dwayne Wade: 1397 points\\n        James Posey: 550 points\\n        ...\\n        ', call_id='call_xdYlGP2lsqDeWdSiOlwOBNiO')], type='ToolCallResultMessage'), TextMessage(source='WebSearchAgent', models_usage=None, content='Tool calls:\\nsearch_web_tool({\"query\":\"Miami Heat highest points player 2006-2007 season\"}) = Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n        Udonis Haslem: 844 points\\n        Dwayne Wade: 1397 points\\n        James Posey: 550 points\\n        ...\\n        ', type='TextMessage'), ToolCallMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=345, completion_tokens=68), content=[FunctionCall(id='call_asjxKtAVGfqrYl0jMpvwCrvV', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='WebSearchAgent'), FunctionCall(id='call_8L91Kizt0KU6RNwUgvNx7S0s', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='WebSearchAgent')], type='ToolCallMessage'), ToolCallResultMessage(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content=\"Error: The tool 'WebSearchAgent' is not available.\", call_id='call_asjxKtAVGfqrYl0jMpvwCrvV'), FunctionExecutionResult(content=\"Error: The tool 'WebSearchAgent' is not available.\", call_id='call_8L91Kizt0KU6RNwUgvNx7S0s')], type='ToolCallResultMessage'), TextMessage(source='DataAnalystAgent', models_usage=None, content='Tool calls:\\nWebSearchAgent({\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}) = Error: The tool \\'WebSearchAgent\\' is not available.\\nWebSearchAgent({\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}) = Error: The tool \\'WebSearchAgent\\' is not available.', type='TextMessage'), ToolCallMessage(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=445, completion_tokens=70), content=[FunctionCall(id='call_imvRJ2jhpPdovBbx8MFjlFVS', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_U30KVmFG1aeXPbqGJjDmJ6iJ', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')], type='ToolCallMessage'), ToolCallResultMessage(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_imvRJ2jhpPdovBbx8MFjlFVS'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_U30KVmFG1aeXPbqGJjDmJ6iJ')], type='ToolCallResultMessage'), TextMessage(source='WebSearchAgent', models_usage=None, content='Tool calls:\\nsearch_web_tool({\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}) = The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\\nsearch_web_tool({\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}) = The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='TextMessage'), ToolCallMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=562, completion_tokens=20), content=[FunctionCall(id='call_CtAnvcbitN0JiwBfiLVzb5Do', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallMessage'), ToolCallResultMessage(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content='85.98130841121495', call_id='call_CtAnvcbitN0JiwBfiLVzb5Do')], type='ToolCallResultMessage'), TextMessage(source='DataAnalystAgent', models_usage=None, content='Tool calls:\\npercentage_change_tool({\"start\":214,\"end\":398}) = 85.98130841121495', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=590, completion_tokens=122), content=\"Summary of Findings:\\n\\n1. Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 season, scoring a total of 1,397 points.\\n2. Dwyane Wade's total rebounds during the 2007-2008 season were 214.\\n3. Dwyane Wade's total rebounds during the 2008-2009 season were 398.\\n4. The percentage change in Dwyane Wade's total rebounds between the 2007-2008 and 2008-2009 seasons was approximately 85.98%.\\n\\nTERMINATE\", type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "task = \"Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\"\n",
    "\n",
    "# Use asyncio.run(...) if you are running this in a script.\n",
    "await Console(team.run_stream(task=task))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, after the Web Search Agent conducts the necessary searches and the Data Analyst Agent completes the necessary calculations, we find that Dwayne Wade was the Miami Heat player with the highest points in the 2006-2007 season, and the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons is 85.98%!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Selector Function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Often times we want better control over the selection process. \n",
    "To this end, we can set the `selector_func` argument with a custom selector function to override the default model-based selection.\n",
    "For instance, we want the Planning Agent to speak immediately after any specialized agent to check the progress.\n",
    "\n",
    "```{note}\n",
    "Returning `None` from the custom selector function will use the default model-based selection.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- user ----------\n",
      "Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?\n",
      "---------- PlanningAgent ----------\n",
      "To address this query, we'll need to break it down into a few specific tasks:\n",
      "\n",
      "1. Web search agent: Identify the Miami Heat player with the highest points in the 2006-2007 NBA season.\n",
      "2. Web search agent: Find the total number of rebounds by this player in the 2007-2008 NBA season.\n",
      "3. Web search agent: Find the total number of rebounds by this player in the 2008-2009 NBA season.\n",
      "4. Data analyst: Calculate the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons.\n",
      "\n",
      "Let's get started with these tasks.\n",
      "[Prompt tokens: 159, Completion tokens: 132]\n",
      "---------- WebSearchAgent ----------\n",
      "[FunctionCall(id='call_TSUHOBKhpHmTNoYeJzwSP5V4', arguments='{\"query\":\"Miami Heat highest points player 2006-2007 season\"}', name='search_web_tool')]\n",
      "[Prompt tokens: 281, Completion tokens: 26]\n",
      "---------- WebSearchAgent ----------\n",
      "[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n        Udonis Haslem: 844 points\\n        Dwayne Wade: 1397 points\\n        James Posey: 550 points\\n        ...\\n        ', call_id='call_TSUHOBKhpHmTNoYeJzwSP5V4')]\n",
      "---------- WebSearchAgent ----------\n",
      "Tool calls:\n",
      "search_web_tool({\"query\":\"Miami Heat highest points player 2006-2007 season\"}) = Here are the total points scored by Miami Heat players in the 2006-2007 season:\n",
      "        Udonis Haslem: 844 points\n",
      "        Dwayne Wade: 1397 points\n",
      "        James Posey: 550 points\n",
      "        ...\n",
      "        \n",
      "---------- PlanningAgent ----------\n",
      "1. Web search agent: Find the total number of rebounds by Dwayne Wade in the 2007-2008 NBA season.\n",
      "2. Web search agent: Find the total number of rebounds by Dwayne Wade in the 2008-2009 NBA season.\n",
      "[Prompt tokens: 382, Completion tokens: 54]\n",
      "---------- DataAnalystAgent ----------\n",
      "[FunctionCall(id='call_BkPBFkpuTG6c3eeoACrrRX7V', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_5LQquT7ZUAAQRf7gvckeTVdQ', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')]\n",
      "[Prompt tokens: 416, Completion tokens: 68]\n",
      "---------- DataAnalystAgent ----------\n",
      "[FunctionExecutionResult(content=\"Error: The tool 'search_web_tool' is not available.\", call_id='call_BkPBFkpuTG6c3eeoACrrRX7V'), FunctionExecutionResult(content=\"Error: The tool 'search_web_tool' is not available.\", call_id='call_5LQquT7ZUAAQRf7gvckeTVdQ')]\n",
      "---------- DataAnalystAgent ----------\n",
      "Tool calls:\n",
      "search_web_tool({\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}) = Error: The tool 'search_web_tool' is not available.\n",
      "search_web_tool({\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}) = Error: The tool 'search_web_tool' is not available.\n",
      "---------- PlanningAgent ----------\n",
      "It seems there was a miscommunication in task assignment. Let me reassess and reassign the tasks correctly.\n",
      "\n",
      "1. Web search agent: Find the total number of rebounds by Dwayne Wade in the 2007-2008 NBA season.\n",
      "2. Web search agent: Find the total number of rebounds by Dwayne Wade in the 2008-2009 NBA season.\n",
      "[Prompt tokens: 525, Completion tokens: 76]\n",
      "---------- WebSearchAgent ----------\n",
      "[FunctionCall(id='call_buIWOtu1dJqPaxJmqMyuRkpj', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_qcnHKdoPsNAzMlPvoBvqmt8n', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')]\n",
      "[Prompt tokens: 599, Completion tokens: 70]\n",
      "---------- WebSearchAgent ----------\n",
      "[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_buIWOtu1dJqPaxJmqMyuRkpj'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_qcnHKdoPsNAzMlPvoBvqmt8n')]\n",
      "---------- WebSearchAgent ----------\n",
      "Tool calls:\n",
      "search_web_tool({\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}) = The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\n",
      "search_web_tool({\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}) = The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.\n",
      "---------- PlanningAgent ----------\n",
      "With this information, we can proceed to calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season to the 2008-2009 season.\n",
      "\n",
      "1. Data analyst: Calculate the percentage change in Dwyane Wade's total rebounds between the 2007-2008 (214 rebounds) and the 2008-2009 (398 rebounds) NBA seasons.\n",
      "[Prompt tokens: 711, Completion tokens: 83]\n",
      "---------- DataAnalystAgent ----------\n",
      "[FunctionCall(id='call_RjbFpLCehz1Nlk5kYmyMUenB', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')]\n",
      "[Prompt tokens: 806, Completion tokens: 20]\n",
      "---------- DataAnalystAgent ----------\n",
      "[FunctionExecutionResult(content='85.98130841121495', call_id='call_RjbFpLCehz1Nlk5kYmyMUenB')]\n",
      "---------- DataAnalystAgent ----------\n",
      "Tool calls:\n",
      "percentage_change_tool({\"start\":214,\"end\":398}) = 85.98130841121495\n",
      "---------- PlanningAgent ----------\n",
      "Based on the data collected, Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 NBA season, scoring a total of 1,397 points. Between the 2007-2008 and 2008-2009 seasons, Dwyane Wade's total rebounds increased from 214 to 398. This represents an approximate 85.98% increase in his total rebounds.\n",
      "\n",
      "TERMINATE\n",
      "[Prompt tokens: 834, Completion tokens: 90]\n",
      "---------- Summary ----------\n",
      "Number of messages: 18\n",
      "Finish reason: Text 'TERMINATE' mentioned\n",
      "Total prompt tokens: 4713\n",
      "Total completion tokens: 619\n",
      "Duration: 11.72 seconds\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Who was the Miami Heat player with the highest points in the 2006-2007 season, and what was the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons?', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=159, completion_tokens=132), content=\"To address this query, we'll need to break it down into a few specific tasks:\\n\\n1. Web search agent: Identify the Miami Heat player with the highest points in the 2006-2007 NBA season.\\n2. Web search agent: Find the total number of rebounds by this player in the 2007-2008 NBA season.\\n3. Web search agent: Find the total number of rebounds by this player in the 2008-2009 NBA season.\\n4. Data analyst: Calculate the percentage change in his total rebounds between the 2007-2008 and 2008-2009 seasons.\\n\\nLet's get started with these tasks.\", type='TextMessage'), ToolCallMessage(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=281, completion_tokens=26), content=[FunctionCall(id='call_TSUHOBKhpHmTNoYeJzwSP5V4', arguments='{\"query\":\"Miami Heat highest points player 2006-2007 season\"}', name='search_web_tool')], type='ToolCallMessage'), ToolCallResultMessage(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n        Udonis Haslem: 844 points\\n        Dwayne Wade: 1397 points\\n        James Posey: 550 points\\n        ...\\n        ', call_id='call_TSUHOBKhpHmTNoYeJzwSP5V4')], type='ToolCallResultMessage'), TextMessage(source='WebSearchAgent', models_usage=None, content='Tool calls:\\nsearch_web_tool({\"query\":\"Miami Heat highest points player 2006-2007 season\"}) = Here are the total points scored by Miami Heat players in the 2006-2007 season:\\n        Udonis Haslem: 844 points\\n        Dwayne Wade: 1397 points\\n        James Posey: 550 points\\n        ...\\n        ', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=382, completion_tokens=54), content='1. Web search agent: Find the total number of rebounds by Dwayne Wade in the 2007-2008 NBA season.\\n2. Web search agent: Find the total number of rebounds by Dwayne Wade in the 2008-2009 NBA season.', type='TextMessage'), ToolCallMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=416, completion_tokens=68), content=[FunctionCall(id='call_BkPBFkpuTG6c3eeoACrrRX7V', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_5LQquT7ZUAAQRf7gvckeTVdQ', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')], type='ToolCallMessage'), ToolCallResultMessage(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content=\"Error: The tool 'search_web_tool' is not available.\", call_id='call_BkPBFkpuTG6c3eeoACrrRX7V'), FunctionExecutionResult(content=\"Error: The tool 'search_web_tool' is not available.\", call_id='call_5LQquT7ZUAAQRf7gvckeTVdQ')], type='ToolCallResultMessage'), TextMessage(source='DataAnalystAgent', models_usage=None, content='Tool calls:\\nsearch_web_tool({\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}) = Error: The tool \\'search_web_tool\\' is not available.\\nsearch_web_tool({\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}) = Error: The tool \\'search_web_tool\\' is not available.', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=525, completion_tokens=76), content='It seems there was a miscommunication in task assignment. Let me reassess and reassign the tasks correctly.\\n\\n1. Web search agent: Find the total number of rebounds by Dwayne Wade in the 2007-2008 NBA season.\\n2. Web search agent: Find the total number of rebounds by Dwayne Wade in the 2008-2009 NBA season.', type='TextMessage'), ToolCallMessage(source='WebSearchAgent', models_usage=RequestUsage(prompt_tokens=599, completion_tokens=70), content=[FunctionCall(id='call_buIWOtu1dJqPaxJmqMyuRkpj', arguments='{\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}', name='search_web_tool'), FunctionCall(id='call_qcnHKdoPsNAzMlPvoBvqmt8n', arguments='{\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}', name='search_web_tool')], type='ToolCallMessage'), ToolCallResultMessage(source='WebSearchAgent', models_usage=None, content=[FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.', call_id='call_buIWOtu1dJqPaxJmqMyuRkpj'), FunctionExecutionResult(content='The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', call_id='call_qcnHKdoPsNAzMlPvoBvqmt8n')], type='ToolCallResultMessage'), TextMessage(source='WebSearchAgent', models_usage=None, content='Tool calls:\\nsearch_web_tool({\"query\": \"Dwyane Wade total rebounds 2007-2008 season\"}) = The number of total rebounds for Dwayne Wade in the Miami Heat season 2007-2008 is 214.\\nsearch_web_tool({\"query\": \"Dwyane Wade total rebounds 2008-2009 season\"}) = The number of total rebounds for Dwayne Wade in the Miami Heat season 2008-2009 is 398.', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=711, completion_tokens=83), content=\"With this information, we can proceed to calculate the percentage change in Dwyane Wade's total rebounds from the 2007-2008 season to the 2008-2009 season.\\n\\n1. Data analyst: Calculate the percentage change in Dwyane Wade's total rebounds between the 2007-2008 (214 rebounds) and the 2008-2009 (398 rebounds) NBA seasons.\", type='TextMessage'), ToolCallMessage(source='DataAnalystAgent', models_usage=RequestUsage(prompt_tokens=806, completion_tokens=20), content=[FunctionCall(id='call_RjbFpLCehz1Nlk5kYmyMUenB', arguments='{\"start\":214,\"end\":398}', name='percentage_change_tool')], type='ToolCallMessage'), ToolCallResultMessage(source='DataAnalystAgent', models_usage=None, content=[FunctionExecutionResult(content='85.98130841121495', call_id='call_RjbFpLCehz1Nlk5kYmyMUenB')], type='ToolCallResultMessage'), TextMessage(source='DataAnalystAgent', models_usage=None, content='Tool calls:\\npercentage_change_tool({\"start\":214,\"end\":398}) = 85.98130841121495', type='TextMessage'), TextMessage(source='PlanningAgent', models_usage=RequestUsage(prompt_tokens=834, completion_tokens=90), content=\"Based on the data collected, Dwyane Wade was the Miami Heat player with the highest points in the 2006-2007 NBA season, scoring a total of 1,397 points. Between the 2007-2008 and 2008-2009 seasons, Dwyane Wade's total rebounds increased from 214 to 398. This represents an approximate 85.98% increase in his total rebounds.\\n\\nTERMINATE\", type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def selector_func(messages: Sequence[AgentMessage]) -> str | None:\n",
    "    if messages[-1].source != planning_agent.name:\n",
    "        return planning_agent.name\n",
    "    return None\n",
    "\n",
    "\n",
    "# Reset the previous team and run the chat again with the selector function.\n",
    "await team.reset()\n",
    "team = SelectorGroupChat(\n",
    "    [planning_agent, web_search_agent, data_analyst_agent],\n",
    "    model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n",
    "    termination_condition=termination,\n",
    "    selector_func=selector_func,\n",
    ")\n",
    "\n",
    "await Console(team.run_stream(task=task))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see from the conversation log that the Planning Agent always speaks immediately after the specialized agents."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}