notebook/agentchat_groupchat_RAG.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Group Chat with Retrieval Augmented Generation\n",
    "\n",
    "AG2 supports conversable agents powered by LLMs, tools, or humans, performing tasks collectively via automated chat. This framework allows tool use and human participation through multi-agent conversation.\n",
    "Please find documentation about this feature [here](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/conversable-agent/).\n",
    "\n",
    "````{=mdx}\n",
    ":::info Requirements\n",
    "Some extra dependencies are needed for this notebook, which can be installed via pip:\n",
    "\n",
    "```bash\n",
    "pip install pyautogen[openai,retrievechat]\n",
    "```\n",
    "\n",
    "For more information, please refer to the [installation guide](https://docs.ag2.ai/docs/user-guide/basic-concepts/installing-ag2).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set your API Endpoint\n",
    "\n",
    "The [`config_list_from_json`](https://docs.ag2.ai/docs/api-reference/autogen/config_list_from_json#config-list-from-json) function loads a list of configurations from an environment variable or a json file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated\n",
    "\n",
    "import autogen\n",
    "from autogen import AssistantAgent\n",
    "from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent\n",
    "\n",
    "llm_config = autogen.LLMConfig.from_json(path=\"OAI_CONFIG_LIST\", timeout=60, temperature=0.8, seed=1234)\n",
    "\n",
    "print(\"LLM models: \", [x.model for x in llm_config.config_list])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{=mdx}\n",
    ":::tip\n",
    "Learn more about configuring LLMs for agents [here](https://docs.ag2.ai/docs/user-guide/basic-concepts/llm-configuration/llm-configuration).\n",
    ":::\n",
    "````\n",
    "\n",
    "## Construct Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def termination_msg(x):\n",
    "    return isinstance(x, dict) and str(x.get(\"content\", \"\"))[-9:].upper() == \"TERMINATE\"\n",
    "\n",
    "\n",
    "boss = autogen.UserProxyAgent(\n",
    "    name=\"Boss\",\n",
    "    is_termination_msg=termination_msg,\n",
    "    human_input_mode=\"NEVER\",\n",
    "    code_execution_config=False,  # we don't want to execute code in this case.\n",
    "    default_auto_reply=\"Reply `TERMINATE` if the task is done.\",\n",
    "    description=\"The boss who ask questions and give tasks.\",\n",
    ")\n",
    "\n",
    "boss_aid = RetrieveUserProxyAgent(\n",
    "    name=\"Boss_Assistant\",\n",
    "    is_termination_msg=termination_msg,\n",
    "    human_input_mode=\"NEVER\",\n",
    "    default_auto_reply=\"Reply `TERMINATE` if the task is done.\",\n",
    "    max_consecutive_auto_reply=3,\n",
    "    retrieve_config={\n",
    "        \"task\": \"code\",\n",
    "        \"docs_path\": \"https://raw.githubusercontent.com/microsoft/FLAML/main/website/docs/Examples/Integrate%20-%20Spark.md\",\n",
    "        \"chunk_token_size\": 1000,\n",
    "        \"model\": llm_config.config_list[0].model,\n",
    "        \"collection_name\": \"groupchat\",\n",
    "        \"get_or_create\": True,\n",
    "    },\n",
    "    code_execution_config=False,  # we don't want to execute code in this case.\n",
    "    description=\"Assistant who has extra content retrieval power for solving difficult problems.\",\n",
    ")\n",
    "\n",
    "coder = AssistantAgent(\n",
    "    name=\"Senior_Python_Engineer\",\n",
    "    is_termination_msg=termination_msg,\n",
    "    system_message=\"You are a senior python engineer, you provide python code to answer questions. Reply `TERMINATE` in the end when everything is done.\",\n",
    "    llm_config=llm_config,\n",
    "    description=\"Senior Python Engineer who can write code to solve problems and answer questions.\",\n",
    ")\n",
    "\n",
    "pm = autogen.AssistantAgent(\n",
    "    name=\"Product_Manager\",\n",
    "    is_termination_msg=termination_msg,\n",
    "    system_message=\"You are a product manager. Reply `TERMINATE` in the end when everything is done.\",\n",
    "    llm_config=llm_config,\n",
    "    description=\"Product Manager who can design and plan the project.\",\n",
    ")\n",
    "\n",
    "reviewer = autogen.AssistantAgent(\n",
    "    name=\"Code_Reviewer\",\n",
    "    is_termination_msg=termination_msg,\n",
    "    system_message=\"You are a code reviewer. Reply `TERMINATE` in the end when everything is done.\",\n",
    "    llm_config=llm_config,\n",
    "    description=\"Code Reviewer who can review the code.\",\n",
    ")\n",
    "\n",
    "PROBLEM = \"How to use spark for parallel training in FLAML? Give me sample code.\"\n",
    "\n",
    "\n",
    "def _reset_agents():\n",
    "    boss.reset()\n",
    "    boss_aid.reset()\n",
    "    coder.reset()\n",
    "    pm.reset()\n",
    "    reviewer.reset()\n",
    "\n",
    "\n",
    "def rag_chat():\n",
    "    _reset_agents()\n",
    "    groupchat = autogen.GroupChat(\n",
    "        agents=[boss_aid, pm, coder, reviewer], messages=[], max_round=12, speaker_selection_method=\"round_robin\"\n",
    "    )\n",
    "    manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)\n",
    "\n",
    "    # Start chatting with boss_aid as this is the user proxy agent.\n",
    "    boss_aid.initiate_chat(\n",
    "        manager,\n",
    "        message=boss_aid.message_generator,\n",
    "        problem=PROBLEM,\n",
    "        n_results=3,\n",
    "    )\n",
    "\n",
    "\n",
    "def norag_chat():\n",
    "    _reset_agents()\n",
    "    groupchat = autogen.GroupChat(\n",
    "        agents=[boss, pm, coder, reviewer],\n",
    "        messages=[],\n",
    "        max_round=12,\n",
    "        speaker_selection_method=\"auto\",\n",
    "        allow_repeat_speaker=False,\n",
    "    )\n",
    "    manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)\n",
    "\n",
    "    # Start chatting with the boss as this is the user proxy agent.\n",
    "    boss.initiate_chat(\n",
    "        manager,\n",
    "        message=PROBLEM,\n",
    "    )\n",
    "\n",
    "\n",
    "def call_rag_chat():\n",
    "    _reset_agents()\n",
    "\n",
    "    # In this case, we will have multiple user proxy agents and we don't initiate the chat\n",
    "    # with RAG user proxy agent.\n",
    "    # In order to use RAG user proxy agent, we need to wrap RAG agents in a function and call\n",
    "    # it from other agents.\n",
    "    def retrieve_content(\n",
    "        message: Annotated[\n",
    "            str,\n",
    "            \"Refined message which keeps the original meaning and can be used to retrieve content for code generation and question answering.\",\n",
    "        ],\n",
    "        n_results: Annotated[int, \"number of results\"] = 3,\n",
    "    ) -> str:\n",
    "        boss_aid.n_results = n_results  # Set the number of results to be retrieved.\n",
    "        _context = {\"problem\": message, \"n_results\": n_results}\n",
    "        ret_msg = boss_aid.message_generator(boss_aid, None, _context)\n",
    "        return ret_msg or message\n",
    "\n",
    "    boss_aid.human_input_mode = \"NEVER\"  # Disable human input for boss_aid since it only retrieves content.\n",
    "\n",
    "    for caller in [pm, coder, reviewer]:\n",
    "        d_retrieve_content = caller.register_for_llm(\n",
    "            description=\"retrieve content for code generation and question answering.\", api_style=\"function\"\n",
    "        )(retrieve_content)\n",
    "\n",
    "    for executor in [boss, pm]:\n",
    "        executor.register_for_execution()(d_retrieve_content)\n",
    "\n",
    "    groupchat = autogen.GroupChat(\n",
    "        agents=[boss, pm, coder, reviewer],\n",
    "        messages=[],\n",
    "        max_round=12,\n",
    "        speaker_selection_method=\"round_robin\",\n",
    "        allow_repeat_speaker=False,\n",
    "    )\n",
    "\n",
    "    manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)\n",
    "\n",
    "    # Start chatting with the boss as this is the user proxy agent.\n",
    "    boss.initiate_chat(\n",
    "        manager,\n",
    "        message=PROBLEM,\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start Chat\n",
    "\n",
    "### UserProxyAgent doesn't get the correct code\n",
    "[FLAML](https://github.com/microsoft/FLAML) was open sourced in 2020, so ChatGPT is familiar with it. However, Spark-related APIs were added in 2022, so they were not in ChatGPT's training data. As a result, we end up with invalid code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norag_chat()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RetrieveUserProxyAgent get the correct code\n",
    "Since RetrieveUserProxyAgent can perform retrieval-augmented generation based on the given documentation file, ChatGPT can generate the correct code for us!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_chat()\n",
    "# type exit to terminate the chat"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Call RetrieveUserProxyAgent while init chat with another user proxy agent\n",
    "Sometimes, there might be a need to use RetrieveUserProxyAgent in group chat without initializing the chat with it. In such scenarios, it becomes essential to create a function that wraps the RAG agents and allows them to be called from other agents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "call_rag_chat()"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Implement and manage a multi-agent chat system using AG2, where AI assistants retrieve information, generate code, and interact collaboratively to solve complex tasks, especially in areas not covered by their training data.",
   "tags": [
    "group chat",
    "orchestration",
    "RAG"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}