{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 多代理辩论\n",
    "\n",
    "多代理辩论是一种多代理设计模式，它模拟多轮交互，在每一轮中，代理之间交换他们的响应，并根据其他代理的响应改进自己的响应。\n",
    "\n",
    "这个示例展示了多代理辩论模式的实现，用于解决来自 [GSM8K 基准测试](https://huggingface.co/datasets/openai/gsm8k)的数学问题。\n",
    "\n",
    "这种模式中有两种类型的代理：求解代理和聚合代理。求解代理按照[Improving Multi-Agent Debate with Sparse Communication Topology](https://arxiv.org/abs/2406.11776)中描述的技术以稀疏方式连接。求解代理负责解决数学问题并与其他代理交换响应。聚合代理负责将数学问题分发给求解代理，等待他们的最终响应，并聚合响应以获得最终答案。\n",
    "\n",
    "该模式的工作方式如下：\n",
    "1. 用户向聚合代理发送数学问题。\n",
    "2. 聚合代理将问题分发给求解代理。\n",
    "3. 每个求解代理处理问题，并向其邻居发布响应。\n",
    "4. 每个求解代理使用来自其邻居的响应来改进其响应，并发布新的响应。\n",
    "5. 重复步骤 4 固定轮数。在最后一轮中，每个求解代理发布最终响应。\n",
    "6. 聚合代理使用多数投票来聚合所有求解代理的最终响应以获得最终答案，并发布答案。\n",
    "\n",
    "我们将使用广播 API，即 {py:meth}`~autogen_core.base.BaseAgent.publish_message`，并且我们将使用主题和订阅来实现通信拓扑。阅读[Topics and Subscriptions](../core-concepts/topic-and-subscription.md)以了解它们的工作原理。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from dataclasses import dataclass\n",
    "from typing import Dict, List\n",
    "\n",
    "from autogen_core.application import SingleThreadedAgentRuntime\n",
    "from autogen_core.base import MessageContext\n",
    "from autogen_core.components import DefaultTopicId, RoutedAgent, TypeSubscription, default_subscription, message_handler\n",
    "from autogen_core.components.models import (\n",
    "    AssistantMessage,\n",
    "    ChatCompletionClient,\n",
    "    LLMMessage,\n",
    "    OpenAIChatCompletionClient,\n",
    "    SystemMessage,\n",
    "    UserMessage,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 消息协议\n",
    "\n",
    "首先，我们定义代理使用的消息。\n",
    "`IntermediateSolverResponse` 是求解代理在每轮中交换的消息，而 `FinalSolverResponse` 是求解代理在最后一轮发布的消息。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Question:\n",
    "    content: str\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Answer:\n",
    "    content: str\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class SolverRequest:\n",
    "    content: str\n",
    "    question: str\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class IntermediateSolverResponse:\n",
    "    content: str\n",
    "    question: str\n",
    "    answer: str\n",
    "    round: int\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class FinalSolverResponse:\n",
    "    answer: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 求解代理\n",
    "\n",
    "求解代理负责解决数学问题并与其他求解代理交换响应。在收到 `SolverRequest` 后，求解代理使用 LLM 生成答案。然后，它根据轮数发布 `IntermediateSolverResponse` 或 `FinalSolverResponse`。\n",
    "\n",
    "求解代理被赋予一个主题类型，用于指示代理应该向其发布中间响应的主题。这个主题由其邻居订阅以接收来自该代理的响应 —— 我们稍后将展示如何实现这一点。\n",
    "\n",
    "我们使用 {py:meth}`~autogen_core.components.default_subscription` 让求解代理订阅默认主题，该主题由聚合代理用来收集求解代理的最终响应。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "@default_subscription\n",
    "class MathSolver(RoutedAgent):\n",
    "    def __init__(self, model_client: ChatCompletionClient, topic_type: str, num_neighbors: int, max_round: int) -> None:\n",
    "        super().__init__(\"A debator.\")\n",
    "        self._topic_type = topic_type\n",
    "        self._model_client = model_client\n",
    "        self._num_neighbors = num_neighbors\n",
    "        self._history: List[LLMMessage] = []\n",
    "        self._buffer: Dict[int, List[IntermediateSolverResponse]] = {}\n",
    "        self._system_messages = [\n",
    "            SystemMessage(\n",
    "                (\n",
    "                    \"You are a helpful assistant with expertise in mathematics and reasoning. \"\n",
    "                    \"Your task is to assist in solving a math reasoning problem by providing \"\n",
    "                    \"a clear and detailed solution. Limit your output within 100 words, \"\n",
    "                    \"and your final answer should be a single numerical number, \"\n",
    "                    \"in the form of {{answer}}, at the end of your response. \"\n",
    "                    \"For example, 'The answer is {{42}}.'\"\n",
    "                )\n",
    "            )\n",
    "        ]\n",
    "        self._round = 0\n",
    "        self._max_round = max_round\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_request(self, message: SolverRequest, ctx: MessageContext) -> None:\n",
    "        # Add the question to the memory.\n",
    "        self._history.append(UserMessage(content=message.content, source=\"user\"))\n",
    "        # Make an inference using the model.\n",
    "        model_result = await self._model_client.create(self._system_messages + self._history)\n",
    "        assert isinstance(model_result.content, str)\n",
    "        # Add the response to the memory.\n",
    "        self._history.append(AssistantMessage(content=model_result.content, source=self.metadata[\"type\"]))\n",
    "        print(f\"{'-'*80}\\nSolver {self.id} round {self._round}:\\n{model_result.content}\")\n",
    "        # Extract the answer from the response.\n",
    "        match = re.search(r\"\\{\\{(\\-?\\d+(\\.\\d+)?)\\}\\}\", model_result.content)\n",
    "        if match is None:\n",
    "            raise ValueError(\"The model response does not contain the answer.\")\n",
    "        answer = match.group(1)\n",
    "        # Increment the counter.\n",
    "        self._round += 1\n",
    "        if self._round == self._max_round:\n",
    "            # If the counter reaches the maximum round, publishes a final response.\n",
    "            await self.publish_message(FinalSolverResponse(answer=answer), topic_id=DefaultTopicId())\n",
    "        else:\n",
    "            # Publish intermediate response to the topic associated with this solver.\n",
    "            await self.publish_message(\n",
    "                IntermediateSolverResponse(\n",
    "                    content=model_result.content,\n",
    "                    question=message.question,\n",
    "                    answer=answer,\n",
    "                    round=self._round,\n",
    "                ),\n",
    "                topic_id=DefaultTopicId(type=self._topic_type),\n",
    "            )\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_response(self, message: IntermediateSolverResponse, ctx: MessageContext) -> None:\n",
    "        # Add neighbor's response to the buffer.\n",
    "        self._buffer.setdefault(message.round, []).append(message)\n",
    "        # Check if all neighbors have responded.\n",
    "        if len(self._buffer[message.round]) == self._num_neighbors:\n",
    "            print(\n",
    "                f\"{'-'*80}\\nSolver {self.id} round {message.round}:\\nReceived all responses from {self._num_neighbors} neighbors.\"\n",
    "            )\n",
    "            # Prepare the prompt for the next question.\n",
    "            prompt = \"These are the solutions to the problem from other agents:\\n\"\n",
    "            for resp in self._buffer[message.round]:\n",
    "                prompt += f\"One agent solution: {resp.content}\\n\"\n",
    "            prompt += (\n",
    "                \"Using the solutions from other agents as additional information, \"\n",
    "                \"can you provide your answer to the math problem? \"\n",
    "                f\"The original math problem is {message.question}. \"\n",
    "                \"Your final answer should be a single numerical number, \"\n",
    "                \"in the form of {{answer}}, at the end of your response.\"\n",
    "            )\n",
    "            # Send the question to the agent itself to solve.\n",
    "            await self.send_message(SolverRequest(content=prompt, question=message.question), self.id)\n",
    "            # Clear the buffer.\n",
    "            self._buffer.pop(message.round)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 聚合代理\n",
    "\n",
    "聚合代理负责处理用户问题并将数学问题分发给求解代理。\n",
    "\n",
    "聚合代理使用 {py:meth}`~autogen_core.components.default_subscription` 订阅默认主题。默认主题用于接收用户问题，接收来自求解代理的最终响应，并将最终答案发布回用户。\n",
    "\n",
    "在更复杂的应用中，当您想要将多代理辩论隔离为一个子组件时，您应该使用 {py:meth}`~autogen_core.components.type_subscription` 为聚合代理-求解代理通信设置特定的主题类型，并让求解代理和聚合代理都发布和订阅该主题类型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "@default_subscription\n",
    "class MathAggregator(RoutedAgent):\n",
    "    def __init__(self, num_solvers: int) -> None:\n",
    "        super().__init__(\"Math Aggregator\")\n",
    "        self._num_solvers = num_solvers\n",
    "        self._buffer: List[FinalSolverResponse] = []\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_question(self, message: Question, ctx: MessageContext) -> None:\n",
    "        print(f\"{'-'*80}\\nAggregator {self.id} received question:\\n{message.content}\")\n",
    "        prompt = (\n",
    "            f\"Can you solve the following math problem?\\n{message.content}\\n\"\n",
    "            \"Explain your reasoning. Your final answer should be a single numerical number, \"\n",
    "            \"in the form of {{answer}}, at the end of your response.\"\n",
    "        )\n",
    "        print(f\"{'-'*80}\\nAggregator {self.id} publishes initial solver request.\")\n",
    "        await self.publish_message(SolverRequest(content=prompt, question=message.content), topic_id=DefaultTopicId())\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_final_solver_response(self, message: FinalSolverResponse, ctx: MessageContext) -> None:\n",
    "        self._buffer.append(message)\n",
    "        if len(self._buffer) == self._num_solvers:\n",
    "            print(f\"{'-'*80}\\nAggregator {self.id} received all final answers from {self._num_solvers} solvers.\")\n",
    "            # Find the majority answer.\n",
    "            answers = [resp.answer for resp in self._buffer]\n",
    "            majority_answer = max(set(answers), key=answers.count)\n",
    "            # Publish the aggregated response.\n",
    "            await self.publish_message(Answer(content=majority_answer), topic_id=DefaultTopicId())\n",
    "            # Clear the responses.\n",
    "            self._buffer.clear()\n",
    "            print(f\"{'-'*80}\\nAggregator {self.id} publishes final answer:\\n{majority_answer}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 设置辩论\n",
    "\n",
    "我们现在将设置一个包含 4 个求解代理和 1 个聚合代理的多代理辩论。求解代理将以稀疏方式连接，如下图所示：\n",
    "\n",
    "```\n",
    "A --- B\n",
    "|     |\n",
    "|     |\n",
    "C --- D\n",
    "```\n",
    "\n",
    "每个求解代理与其他两个求解代理相连。例如，代理 A 与代理 B 和 C 相连。\n",
    "\n",
    "让我们首先创建一个运行时并注册代理类型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AgentType(type='MathAggregator')"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "runtime = SingleThreadedAgentRuntime()\n",
    "await MathSolver.register(\n",
    "    runtime,\n",
    "    \"MathSolverA\",\n",
    "    lambda: MathSolver(\n",
    "        model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n",
    "        topic_type=\"MathSolverA\",\n",
    "        num_neighbors=2,\n",
    "        max_round=3,\n",
    "    ),\n",
    ")\n",
    "await MathSolver.register(\n",
    "    runtime,\n",
    "    \"MathSolverB\",\n",
    "    lambda: MathSolver(\n",
    "        model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n",
    "        topic_type=\"MathSolverB\",\n",
    "        num_neighbors=2,\n",
    "        max_round=3,\n",
    "    ),\n",
    ")\n",
    "await MathSolver.register(\n",
    "    runtime,\n",
    "    \"MathSolverC\",\n",
    "    lambda: MathSolver(\n",
    "        model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n",
    "        topic_type=\"MathSolverC\",\n",
    "        num_neighbors=2,\n",
    "        max_round=3,\n",
    "    ),\n",
    ")\n",
    "await MathSolver.register(\n",
    "    runtime,\n",
    "    \"MathSolverD\",\n",
    "    lambda: MathSolver(\n",
    "        model_client=OpenAIChatCompletionClient(model=\"gpt-4o-mini\"),\n",
    "        topic_type=\"MathSolverD\",\n",
    "        num_neighbors=2,\n",
    "        max_round=3,\n",
    "    ),\n",
    ")\n",
    "await MathAggregator.register(runtime, \"MathAggregator\", lambda: MathAggregator(num_solvers=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在我们将使用 {py:class}`~autogen_core.components.TypeSubscription` 创建求解代理拓扑，它将每个求解代理的发布主题类型映射到其邻居的代理类型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Subscriptions for topic published to by MathSolverA.\n",
    "await runtime.add_subscription(TypeSubscription(\"MathSolverA\", \"MathSolverD\"))\n",
    "await runtime.add_subscription(TypeSubscription(\"MathSolverA\", \"MathSolverB\"))\n",
    "\n",
    "# Subscriptions for topic published to by MathSolverB.\n",
    "await runtime.add_subscription(TypeSubscription(\"MathSolverB\", \"MathSolverA\"))\n",
    "await runtime.add_subscription(TypeSubscription(\"MathSolverB\", \"MathSolverC\"))\n",
    "\n",
    "# Subscriptions for topic published to by MathSolverC.\n",
    "await runtime.add_subscription(TypeSubscription(\"MathSolverC\", \"MathSolverB\"))\n",
    "await runtime.add_subscription(TypeSubscription(\"MathSolverC\", \"MathSolverD\"))\n",
    "\n",
    "# Subscriptions for topic published to by MathSolverD.\n",
    "await runtime.add_subscription(TypeSubscription(\"MathSolverD\", \"MathSolverC\"))\n",
    "await runtime.add_subscription(TypeSubscription(\"MathSolverD\", \"MathSolverA\"))\n",
    "\n",
    "# All solvers and the aggregator subscribe to the default topic."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 解决数学问题\n",
    "\n",
    "现在让我们运行辩论来解决一个数学问题。我们向默认主题发布一个 `SolverRequest`，聚合代理将开始辩论。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "Aggregator MathAggregator:default received question:\n",
      "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?\n",
      "--------------------------------------------------------------------------------\n",
      "Aggregator MathAggregator:default publishes initial solver request.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverC:default round 0:\n",
      "In April, Natalia sold 48 clips. In May, she sold half as many, which is 48 / 2 = 24 clips. To find the total number of clips sold in April and May, we add the amounts: 48 (April) + 24 (May) = 72 clips. \n",
      "\n",
      "Thus, the total number of clips sold by Natalia is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverB:default round 0:\n",
      "In April, Natalia sold 48 clips. In May, she sold half as many clips, which is 48 / 2 = 24 clips. To find the total clips sold in April and May, we add both amounts: \n",
      "\n",
      "48 (April) + 24 (May) = 72.\n",
      "\n",
      "Thus, the total number of clips sold altogether is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverD:default round 0:\n",
      "Natalia sold 48 clips in April. In May, she sold half as many, which is \\( \\frac{48}{2} = 24 \\) clips. To find the total clips sold in both months, we add the clips sold in April and May together:\n",
      "\n",
      "\\[ 48 + 24 = 72 \\]\n",
      "\n",
      "Thus, Natalia sold a total of 72 clips.\n",
      "\n",
      "The answer is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverC:default round 1:\n",
      "Received all responses from 2 neighbors.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverA:default round 1:\n",
      "Received all responses from 2 neighbors.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverA:default round 0:\n",
      "In April, Natalia sold clips to 48 friends. In May, she sold half as many, which is calculated as follows:\n",
      "\n",
      "Half of 48 is \\( 48 \\div 2 = 24 \\).\n",
      "\n",
      "Now, to find the total clips sold in April and May, we add the totals from both months:\n",
      "\n",
      "\\( 48 + 24 = 72 \\).\n",
      "\n",
      "Thus, the total number of clips Natalia sold altogether in April and May is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverD:default round 1:\n",
      "Received all responses from 2 neighbors.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverB:default round 1:\n",
      "Received all responses from 2 neighbors.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverC:default round 1:\n",
      "In April, Natalia sold 48 clips. In May, she sold half as many, which is 48 / 2 = 24 clips. The total number of clips sold in April and May is calculated by adding the two amounts: 48 (April) + 24 (May) = 72 clips. \n",
      "\n",
      "Therefore, the answer is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverA:default round 1:\n",
      "In April, Natalia sold 48 clips. In May, she sold half of that amount, which is 48 / 2 = 24 clips. To find the total clips sold in both months, we sum the clips from April and May: \n",
      "\n",
      "48 (April) + 24 (May) = 72.\n",
      "\n",
      "Thus, Natalia sold a total of {{72}} clips. \n",
      "\n",
      "The answer is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverD:default round 2:\n",
      "Received all responses from 2 neighbors.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverB:default round 2:\n",
      "Received all responses from 2 neighbors.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverD:default round 1:\n",
      "Natalia sold 48 clips in April. In May, she sold half of that, which is \\( 48 \\div 2 = 24 \\) clips. To find the total clips sold, we add the clips sold in both months:\n",
      "\n",
      "\\[ 48 + 24 = 72 \\]\n",
      "\n",
      "Therefore, the total number of clips sold by Natalia is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverB:default round 1:\n",
      "In April, Natalia sold 48 clips. In May, she sold half that amount, which is 48 / 2 = 24 clips. To find the total clips sold in both months, we add the amounts: \n",
      "\n",
      "48 (April) + 24 (May) = 72.\n",
      "\n",
      "Therefore, the total number of clips sold altogether by Natalia is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverA:default round 2:\n",
      "Received all responses from 2 neighbors.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverC:default round 2:\n",
      "Received all responses from 2 neighbors.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverA:default round 2:\n",
      "In April, Natalia sold 48 clips. In May, she sold half of that amount, which is \\( 48 \\div 2 = 24 \\) clips. To find the total clips sold in both months, we add the amounts from April and May:\n",
      "\n",
      "\\( 48 + 24 = 72 \\).\n",
      "\n",
      "Thus, the total number of clips sold by Natalia is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverC:default round 2:\n",
      "In April, Natalia sold 48 clips. In May, she sold half of that amount, which is \\( 48 \\div 2 = 24 \\) clips. To find the total number of clips sold in both months, we add the clips sold in April and May: \n",
      "\n",
      "48 (April) + 24 (May) = 72. \n",
      "\n",
      "Thus, the total number of clips sold altogether by Natalia is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverB:default round 2:\n",
      "In April, Natalia sold 48 clips. In May, she sold half as many, calculated as \\( 48 \\div 2 = 24 \\) clips. To find the total clips sold over both months, we sum the totals: \n",
      "\n",
      "\\( 48 (April) + 24 (May) = 72 \\).\n",
      "\n",
      "Therefore, the total number of clips Natalia sold is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Solver MathSolverD:default round 2:\n",
      "To solve the problem, we know that Natalia sold 48 clips in April. In May, she sold half that amount, which is calculated as \\( 48 \\div 2 = 24 \\) clips. To find the total number of clips sold over both months, we add the two amounts together:\n",
      "\n",
      "\\[ 48 + 24 = 72 \\]\n",
      "\n",
      "Thus, the total number of clips sold by Natalia is {{72}}.\n",
      "--------------------------------------------------------------------------------\n",
      "Aggregator MathAggregator:default received all final answers from 4 solvers.\n",
      "--------------------------------------------------------------------------------\n",
      "Aggregator MathAggregator:default publishes final answer:\n",
      "72\n"
     ]
    }
   ],
   "source": [
    "question = \"Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?\"\n",
    "runtime.start()\n",
    "await runtime.publish_message(Question(content=question), DefaultTopicId())\n",
    "await runtime.stop_when_idle()"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}