{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z4UWyDywJ6tj"
      },
      "source": [
        "# AG2 (formerly Autogen) Multi-Agents Example on Vertex AI Agent Engine\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_ag2_on_agent_engine.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fagent-engine%2Ftutorial_ag2_on_agent_engine.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/agent-engine/tutorial_ag2_on_agent_engine.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_ag2_on_agent_engine.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_ag2_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_ag2_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_ag2_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_ag2_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/agent-engine/tutorial_ag2_on_agent_engine.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>\n",
        "\n",
        "\n",
        "<br>\n",
        "<br>\n",
        "<br>\n",
        "<br>\n",
        "\n",
        "| | |\n",
        "|-|-|\n",
        "| Author(s) | [Emmanuel Awa](https://github.com/awaemmanuel)|\n",
        "| Reviewer(s) | Turan Bulmus, [Rajesh Thallam](https://github.com/rthallam), Polong Lin|\n",
        "| Last updated | 2025 03 16: Update to latest Agent Engine SDK |\n",
        "| | 2025 02 20: Update to latest AG2 SDK and report summarization using LLM |\n",
        "| | 2025 01 16: Updated from `autogen[gemini]` to `ag2`  |\n",
        "| | 2024 12 20: Review |\n",
        "| | 2024 12 19: Initial Publication|\n",
        "\n",
        "This notebook showcases how to deploy [AG2](https://ag2.ai/) on Vertex AI through [Agent Engine API](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/reasoning-engine). AG2 (formerly AutoGen) is an open-source framework for building AI agents that can collaborate to solve tasks. It simplifies agent development, supports LLMs and tools, and enables various workflows, making it suitable for building multiagent applications on Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VyKi7eTEgVwk"
      },
      "source": [
        "## Prerequisites\n",
        "In other to run this notebook, you will need to set up your Google Cloud project and enable the Vertex AI API.   \n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager).\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "1. [Enable the Service Usage API](https://console.cloud.google.com/apis/library/serviceusage.googleapis.com)\n",
        "1. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "1. [Enable Logs Explorer](https://console.cloud.google.com/logs/query;query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7NX_Q-pOqMRi"
      },
      "source": [
        "## Step 1: Environment Setup: Google Cloud and AG2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8yfWpe5Nctnb"
      },
      "source": [
        "###  Authentication  \n",
        "\n",
        "If you're using Colab, run the code in the next cell. Follow the pop ups and authenticate with an account that has access to your Google Cloud [project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#identifying_projects).\n",
        "\n",
        "If you're running this notebook somewhere besides Colab, make sure your environment has the right Google Cloud access. If that's a new concept to you, consider looking into [Application Default Credentials for your local environment](https://cloud.google.com/docs/authentication/provide-credentials-adc#local-dev) and [initializing the Google Cloud CLI](https://cloud.google.com/docs/authentication/gcloud).\n",
        "\n",
        "In many cases, running `gcloud auth application-default login` in a shell on the machine running the notebook kernel is sufficient.\n",
        "\n",
        "More authentication options are discussed [here](https://cloud.google.com/docs/authentication)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rPZA1RYSj0Vs"
      },
      "outputs": [],
      "source": [
        "# @title Authenticate - Google Colab Only\n",
        "import sys\n",
        "\n",
        "from google.colab import auth\n",
        "\n",
        "if \"google.colab.auth\" in sys.modules:\n",
        "    auth.authenticate_user()\n",
        "    print(\"Authenticated\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jQs3DQ4DxB_k"
      },
      "source": [
        "###  Install Dependencies and Import Packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jv4YSAecekxA"
      },
      "outputs": [],
      "source": [
        "# @title Install AG2, Google Cloud SDK and Dependencies\n",
        "%pip install -qqq ag2[gemini]==0.8.3 dask[dataframe]==2024.12.1 google-cloud-aiplatform[agent_engines,ag2]==1.83.0\n",
        "%pip show ag2 google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "txTCLHoghI2F"
      },
      "outputs": [],
      "source": [
        "# @title Import relevant packages\n",
        "import pprint\n",
        "\n",
        "from IPython.display import Markdown\n",
        "from autogen import AssistantAgent, GroupChat, GroupChatManager, UserProxyAgent\n",
        "from autogen.cache import Cache\n",
        "from autogen.code_utils import content_str\n",
        "import vertexai\n",
        "from vertexai.agent_engines import AgentEngine, Queryable\n",
        "\n",
        "print(f\"Required libraries successfully imported.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6M7rdk_jw2bq"
      },
      "source": [
        "Next, we need to first install AG2 with Gemini features."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UmnZKKAnxHz6"
      },
      "source": [
        "### Configuration Parameters\n",
        "\n",
        "Before proceeding, you need to define the following parameters:\n",
        "\n",
        "1.  **`MODEL`**: Specifies the Gemini model to utilize for the multi-agent conversation. Choose from available Gemini model options outlined [here](https://ai.google.dev/gemini-api/docs/models/gemini).\n",
        "2.  **`PROJECT_ID`**: Identifies the Google Cloud project where the executions will occur.\n",
        "3.  **`STAGING_BUCKET`**: Designates a Google Cloud Storage bucket that you own and have write access to. This bucket will be used for temporary storage during the execution process. Bucket name must begin with `gs://`\n",
        "4.  **`CACHING_SEED`**:  Provides an integer value for best-effort deterministic sampling. Note that determinism is not guaranteed.\n",
        "5.  **`LOCATION`**:  Indicates the Google Cloud region for the deployment. Currently, Agent Engine is only supported in `us-central1`.\n",
        "\n",
        "These parameters will be employed to configure the `AG2` framework and set up the remote deployment environment for the Agent Engine."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-eVuRe2efbnj"
      },
      "outputs": [],
      "source": [
        "# @title Set Configuration Parameters\n",
        "MODEL_NAME = \"gemini-2.0-flash-001\"  # @param {'type': 'string'}\n",
        "PROJECT_ID = \"<YOUR_PROJECT_ID>\"  # @param {'type': 'string'}\n",
        "STAGING_BUCKET = \"gs://<YOUR_GCS_BUCKET>\"  # @param {'type': 'string'}\n",
        "CACHING_SEED = 42  # @param {'type': 'integer'}\n",
        "LOCATION = \"us-central1\"\n",
        "\n",
        "STAGING_BUCKET = (\n",
        "    STAGING_BUCKET if STAGING_BUCKET.startswith(\"gs://\") else f\"gs://{STAGING_BUCKET}\"\n",
        ")\n",
        "\n",
        "config_list = [\n",
        "    {\n",
        "        \"model\": MODEL_NAME,\n",
        "        \"project_id\": PROJECT_ID,\n",
        "        \"location\": LOCATION,\n",
        "        \"api_type\": \"google\",\n",
        "    }\n",
        "]\n",
        "seed = CACHING_SEED\n",
        "all_config = {\n",
        "    \"gemini_config_list\": config_list,\n",
        "    \"caching_seed\": seed,\n",
        "    \"staging_bucket\": STAGING_BUCKET,\n",
        "}\n",
        "print(\"Configuration:\")\n",
        "pprint.pprint(all_config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "30HHGyoChnay"
      },
      "outputs": [],
      "source": [
        "# @title Initialize Vertex AI\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)\n",
        "print(f\"Vertex AI successfully initialized\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TOPhtR7tfYls"
      },
      "source": [
        "## Step 2: An AG2 Example of Multi-Agent Conversations\n",
        "\n",
        "AG2 supports multi-agent collaboration through a conversational pattern called **group chat**. This general pattern shares characteristics with known multi-agent patterns like **Router Agent (one-by-one)**, where a central agent delegates tasks sequentially, and **Dynamic (all-to-all)**, where agents communicate freely.\n",
        "\n",
        "However, AG2's `GroupChat` offers a distinct approach. A `GroupChatManager` orchestrates the conversation by selecting a speaker, broadcasting their message, and repeating the process. While not strictly sequential like a Router, the manager guides the flow depending on the strategy the user chooses.\n",
        "\n",
        "Unlike a purely Dynamic system, the manager provides a degree of control, preventing chaotic interactions. Agent selection strategies include: `round_robin`, `random`, `manual`, and the default `auto` (based on the manager's llm configuration). We have chosen the `auto` mode to allow Gemini's reasoning ability decide who the next speaker should be.\n",
        "\n",
        "This example showcases a research-focused group chat with five agents:\n",
        "\n",
        "*   **Researcher:** Conducts research and identifies key information.\n",
        "*   **Data Analyst:** Analyzes data and extracts insights.\n",
        "*   **Literature Reviewer:** Synthesizes research into a literature review.\n",
        "*   **Fact Checker:** Verifies information accuracy.\n",
        "*   **User Proxy:** Represents the user, initiates the conversation, and can execute code.\n",
        "\n",
        "These agents reside in a `GroupChat` orchestrated by a `GroupChatManager` using the `auto` strategy. The `User_proxy` initiates the collaborative research workflow.\n",
        "\n",
        "The diagram below shows the `GroupChatManager` workflow.\n",
        "\n",
        "> **NOTE:** If the image is not legible, please right click and view in new tab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DThm2rfjZCY_"
      },
      "source": [
        "![Ag2](https://storage.googleapis.com/github-repo/generative-ai/gemini/agent-engine/ag2/ag2.jpeg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CbqXfcfi5Zsm"
      },
      "source": [
        "**Defining the Research Topic**\n",
        "\n",
        "This section allows you to specify the topic of interest for the multi-agent research team. By default, the agents will focus on \"Large Language Models.\" However, you can customize the research topic by modifying the `topic` below:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_-rGbQxQ5J8L"
      },
      "outputs": [],
      "source": [
        "research_topic = \"Large Language Models\"  # @param [\"Large Language Models\", \"Real Estate\", \"Agentic Workflows\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hvqc_EdH6cTX"
      },
      "source": [
        "The specialized agents (researcher, fact-checker, data analysts and literature reviewer) utilize this topic as their primary focus for information gathering and analysis.  By defining a clear research topic, you ensure that the multi-agent conversation remains relevant and productive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RTJBQ5_OngMh"
      },
      "outputs": [],
      "source": [
        "# @title Define Agents\n",
        "\n",
        "# Define the Researcher agent. This agent conducts thorough research on the given topic,\n",
        "# identifying key trends, publications, experts, and relevant data sources.\n",
        "researcher = AssistantAgent(\n",
        "    name=\"Researcher\",\n",
        "    llm_config={\"config_list\": config_list, \"seed\": seed},\n",
        "    description=\"Conducts thorough research on the given topic.\",\n",
        "    system_message=f\"\"\"You are a meticulous researcher. Your goal is to conduct thorough research, identifying key trends, publications, experts, and relevant data sources. You should provide accurate and well-cited information.\"\"\",\n",
        ")\n",
        "\n",
        "# Define the Literature Reviewer agent. This agent synthesizes research findings into a\n",
        "# coherent and well-structured literature review, identifying key themes, connections, and gaps in the existing research.\n",
        "literature_reviewer = AssistantAgent(\n",
        "    name=\"Literature_Reviewer\",\n",
        "    llm_config={\"config_list\": config_list, \"seed\": seed},\n",
        "    description=\"Synthesizes research findings into a coherent literature review.\",\n",
        "    system_message=f\"\"\"You are a skilled literature reviewer. Your goal is to synthesize the provided research findings into a coherent and well-structured literature review. You should identify key themes, connections, and gaps in the existing research.\"\"\",\n",
        ")\n",
        "\n",
        "# Define the Fact Checker agent. This agent verifies the accuracy and validity of information,\n",
        "# checking sources, identifying potential biases, and ensuring consistency with established knowledge.\n",
        "fact_checker = AssistantAgent(\n",
        "    name=\"Fact_Checker\",\n",
        "    llm_config={\"config_list\": config_list, \"seed\": seed},\n",
        "    description=\"Verifies the accuracy and validity of information.\",\n",
        "    system_message=\"\"\"You are a meticulous fact-checker. Your goal is to verify the accuracy and validity of any information presented to you. You should check sources, identify potential biases, and ensure the information is consistent with established knowledge.\"\"\",\n",
        ")\n",
        "\n",
        "# Define the Data Analyst agent. This agent analyzes data related to the research topic and\n",
        "# extracts meaningful insights, identifying trends, patterns, and correlations.\n",
        "data_analyst = AssistantAgent(\n",
        "    name=\"Data_Analyst\",\n",
        "    llm_config={\"config_list\": config_list, \"seed\": seed},\n",
        "    description=\"Analyzes data and extracts meaningful insights.\",\n",
        "    system_message=f\"\"\"You are a skilled data analyst. Your goal is to analyze data and extract meaningful insights. You should identify trends, patterns, and correlations in the data and present your findings in a clear and concise manner.\"\"\",\n",
        ")\n",
        "\n",
        "# Define the User Proxy agent. This agent represents the user and can run code to facilitate tasks.\n",
        "user_proxy = UserProxyAgent(\n",
        "    name=\"User_proxy\",\n",
        "    code_execution_config={\n",
        "        \"last_n_messages\": 20,\n",
        "        \"work_dir\": \"coding\",\n",
        "        \"use_docker\": False,\n",
        "    },\n",
        "    human_input_mode=\"TERMINATE\",\n",
        "    is_termination_msg=lambda x: content_str(x.get(\"content\")).find(\"TERMINATE\") >= 0,\n",
        "    description=\"I represent the user, and can run code.\",\n",
        ")\n",
        "\n",
        "# Create the GroupChat, a conversable pattern, that facilitates the collaboration on tasks requiring interaction among multiple agents.\n",
        "# All agents contribute to a single conversation thread and share the same context. We have set the max turns of conversation to 12\n",
        "groupchat = GroupChat(\n",
        "    agents=[user_proxy, researcher, literature_reviewer, fact_checker, data_analyst],\n",
        "    messages=[],\n",
        "    max_round=12,\n",
        ")\n",
        "\n",
        "# Create the GroupChatManager, a specialized agent, to orchestrate the group chat. A GroupChatManager selects an agent to speak,\n",
        "# broadcasts the message to all other agents, and repeats this process until the conversation stops or maximum number of turns is reached.\n",
        "# The default selection strategy is \"auto\", where the manager's LLM, in our case Gemini, chooses the next speaker.\n",
        "manager = GroupChatManager(\n",
        "    groupchat=groupchat,\n",
        "    llm_config={\"config_list\": config_list, \"seed\": seed},\n",
        "    is_termination_msg=lambda x: content_str(x.get(\"content\")).find(\"TERMINATE\") >= 0,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fXcoq_nrBLgu"
      },
      "outputs": [],
      "source": [
        "# @title Start Multi-agent Conversations\n",
        "msg = f\"\"\"Conduct research on the latest advancements in {research_topic}, analyze relevant data, and produce a comprehensive literature review with verified facts.\"\"\"\n",
        "custom_summary_prompt = f\"Generate a final report on {research_topic} from the chat history. Just start the report immediately without adding any additional commentary before or after the report.\"\n",
        "with Cache.disk() as cache:  # enable caching to reduce inference cost\n",
        "    result = user_proxy.initiate_chat(\n",
        "        manager,\n",
        "        message=msg,\n",
        "        summary_method=\"reflection_with_llm\",\n",
        "        summary_args={\"summary_prompt\": custom_summary_prompt},\n",
        "        cache=cache,\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QpaKbuqs_ZwQ"
      },
      "outputs": [],
      "source": [
        "# @title Inspect Report Generated By Agents\n",
        "summary = result.summary\n",
        "final_report = content_str(\n",
        "    summary.get(\"content\") if isinstance(summary, dict) else summary\n",
        ")\n",
        "display(Markdown(final_report))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w1YCBTrbZtlg"
      },
      "source": [
        "## Step 3. Deploy agents to Agent Engine"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uXgl1m3cxPW3"
      },
      "source": [
        "### Deploying The AG2 Multi-Agent Conversations Example On Vertex AI Agent Engine.\n",
        "\n",
        "1. Create a class `ResearchApp`, that implements a agent engine `Queryable` class. This class will wrap the agents from cell above.\n",
        "\n",
        "2. The syntax used to create a register a agent engine resource needs to implement `set_up` and `query` methods. Here is an example of a Agent Engine [advanced app configuration](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/reasoning-engine#deploy_an_advanced_app_configuration) that show this.\n",
        "\n",
        "\n",
        "> **NOTE:** All agents and boostrapping code must be initialized in the `set_up` method and not the `__init__` method. Multi-agent conversation and user proxy code will go into the `query` method.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Gr24G_qzFhZ9"
      },
      "outputs": [],
      "source": [
        "# @title Define AG2 Agents As a Agent Engine QueryAble\n",
        "\n",
        "\n",
        "class ResearchApp(Queryable):\n",
        "    \"\"\"This class defines a research application using AG2. It facilitates multi-agent conversations for research tasks.\"\"\"\n",
        "\n",
        "    base_question = \"Conduct research on the latest advancements in {topic}, analyze relevant data, and produce a comprehensive literature review with verified facts.\"\n",
        "    custom_summary_prompt = f\"Generate a final report on {research_topic} from the chat history. Just start the report immediately without adding any additional commentary before or after the report.\"\n",
        "\n",
        "    def __init__(self, project_id: str, location: str):\n",
        "        \"\"\"\n",
        "        Initializes the ResearchApp instance.\n",
        "\n",
        "        Args:\n",
        "            project_id (str): The project ID for AG2.\n",
        "            location (str): The location of the AG2 instance.\n",
        "        \"\"\"\n",
        "        self.project_id = project_id\n",
        "        self.location = location\n",
        "        self.set_up_complete = False\n",
        "\n",
        "    def set_up(self):\n",
        "        \"\"\"Sets up the research agents and conversation environment.\"\"\"\n",
        "        self.researcher = AssistantAgent(\n",
        "            name=\"Researcher\",\n",
        "            llm_config={\n",
        "                \"config_list\": all_config[\"gemini_config_list\"],\n",
        "                \"seed\": all_config[\"caching_seed\"],\n",
        "            },\n",
        "            description=\"Conducts thorough research on the given topic.\",\n",
        "            system_message=\"\"\"You are a meticulous researcher. Your goal is to conduct thorough research, identifying key trends, publications, experts, and relevant data sources. You should provide accurate and well-cited information.\"\"\",\n",
        "        )\n",
        "\n",
        "        self.literature_reviewer = AssistantAgent(\n",
        "            name=\"Literature_Reviewer\",\n",
        "            llm_config={\n",
        "                \"config_list\": all_config[\"gemini_config_list\"],\n",
        "                \"seed\": all_config[\"caching_seed\"],\n",
        "            },\n",
        "            description=\"Synthesizes research findings into a coherent literature review.\",\n",
        "            system_message=\"\"\"You are a skilled literature reviewer. Your goal is to synthesize the provided research findings into a coherent and well-structured literature review. You should identify key themes, connections, and gaps in the existing research.\"\"\",\n",
        "        )\n",
        "\n",
        "        self.fact_checker = AssistantAgent(\n",
        "            name=\"Fact_Checker\",\n",
        "            llm_config={\n",
        "                \"config_list\": all_config[\"gemini_config_list\"],\n",
        "                \"seed\": all_config[\"caching_seed\"],\n",
        "            },\n",
        "            description=\"Verifies the accuracy and validity of information.\",\n",
        "            system_message=\"\"\"You are a meticulous fact-checker. Your goal is to verify the accuracy and validity of any information presented to you. You should check sources, identify potential biases, and ensure the information is consistent with established knowledge.\"\"\",\n",
        "        )\n",
        "\n",
        "        self.data_analyst = AssistantAgent(\n",
        "            name=\"Data_Analyst\",\n",
        "            llm_config={\n",
        "                \"config_list\": all_config[\"gemini_config_list\"],\n",
        "                \"seed\": all_config[\"caching_seed\"],\n",
        "            },\n",
        "            description=\"Analyzes data and extracts meaningful insights.\",\n",
        "            system_message=\"\"\"You are a skilled data analyst. Your goal is to analyze data and extract meaningful insights. You should identify trends, patterns, and correlations in the data and present your findings in a clear and concise manner.\"\"\",\n",
        "        )\n",
        "\n",
        "        self.user_proxy = UserProxyAgent(\n",
        "            name=\"User_proxy\",\n",
        "            code_execution_config={\n",
        "                \"last_n_messages\": 20,\n",
        "                \"work_dir\": \"coding\",\n",
        "                \"use_docker\": False,\n",
        "            },\n",
        "            human_input_mode=\"NEVER\",\n",
        "            is_termination_msg=lambda x: content_str(x.get(\"content\")).find(\"TERMINATE\")\n",
        "            >= 0,\n",
        "            description=\"I represent the user, and can run code.\",\n",
        "        )\n",
        "\n",
        "        self.groupchat = GroupChat(\n",
        "            agents=[\n",
        "                self.user_proxy,\n",
        "                self.researcher,\n",
        "                self.literature_reviewer,\n",
        "                self.fact_checker,\n",
        "                self.data_analyst,\n",
        "            ],\n",
        "            messages=[],\n",
        "            max_round=12,\n",
        "        )\n",
        "        self.manager = GroupChatManager(\n",
        "            groupchat=self.groupchat,\n",
        "            llm_config={\"config_list\": config_list, \"seed\": seed},\n",
        "            is_termination_msg=lambda x: content_str(x.get(\"content\")).find(\"TERMINATE\")\n",
        "            >= 0,\n",
        "        )\n",
        "        self.set_up_complete = True\n",
        "\n",
        "    def query(self, topic: str = \"Large Language Models\"):\n",
        "        \"\"\"\n",
        "        Queries the multi-agent research system.\n",
        "\n",
        "        Args:\n",
        "            topic (str): The research topic. Defaults to \"Large Language Models\".\n",
        "\n",
        "        Returns:\n",
        "            str: A summary of the research conducted by the multi-agent system.\n",
        "        \"\"\"\n",
        "        if not self.set_up_complete:\n",
        "            self.set_up()\n",
        "        question = self.base_question.format(topic=topic)\n",
        "        with Cache.disk() as cache:\n",
        "            result = self.user_proxy.initiate_chat(\n",
        "                self.manager,\n",
        "                message=question,\n",
        "                summary_method=\"reflection_with_llm\",\n",
        "                summary_args={\"summary_prompt\": custom_summary_prompt},\n",
        "                cache=cache,\n",
        "            )\n",
        "        summary = result.summary\n",
        "        return content_str(\n",
        "            summary.get(\"content\") if isinstance(summary, dict) else summary\n",
        "        )\n",
        "\n",
        "\n",
        "print(f\"Research App with multi-agents created.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Kwmo5QJx8Dv"
      },
      "source": [
        "### Test\n",
        "\n",
        "This section demonstrates how to test the Agent Engine locally before deploying it to Vertex AI. By running the `ResearchApp` locally, you can observe a multi-agent conversation flow similar to the previous example. This local testing helps ensure that the agents and their interactions are functioning as expected before deploying to the cloud."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z89kARQryAfb"
      },
      "outputs": [],
      "source": [
        "# @title Testing - Local Development\n",
        "local_engine = ResearchApp(project_id=PROJECT_ID, location=LOCATION)\n",
        "local_engine.set_up()\n",
        "report = local_engine.query(topic=research_topic)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "uM5QZ77lfo0Q"
      },
      "outputs": [],
      "source": [
        "# @title Inspect Report\n",
        "display(Markdown(report))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eOk6mEfYxl9n"
      },
      "source": [
        "### Setup and Deploy to Vertex AI Agent Engine\n",
        "\n",
        "This section guides you through deploying the AG2 multi-agent application to Vertex AI Agent Engine. Before proceeding with the deployment, we'll ensure there are no existing deployments to avoid conflicts. This involves listing and deleting any previous Agent Engine deployments to start with a clean slate. Subsequently, we'll utilize the `AgentEngine.create` method to deploy the application, specifying dependencies and configurations.\n",
        "\n",
        "> **NOTE:** Please be patient as the deployment is a long running operation. Once successful, the cell will return the resource URI of deployment that can be used for subsequent executions.  \n",
        ">**Example:**\n",
        "> ` agent_engine = vertexai.agent_engines.get('projects/<YOUR_PROJECT_ID_NUMBER>/locations/us-central1/reasoningEngines/<DEPLOYED_AGENT_ENGINE_ID>')`\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "BpXwgjnrXLto"
      },
      "outputs": [],
      "source": [
        "# @title Clean up deployment if needed\n",
        "\n",
        "\n",
        "def cleanup_deployments(dry_run: bool = True):\n",
        "    agent_engine_list = AgentEngine.list()\n",
        "    if agent_engine_list:\n",
        "        print(f\"Found {len(agent_engine_list)} existing deployments:\")\n",
        "        for re in agent_engine_list:\n",
        "            print(f\"Deleting {re.display_name}\")\n",
        "            print(f\"Resource URI: {re.resource_name}\")\n",
        "            if not dry_run:\n",
        "                re.delete()\n",
        "            else:\n",
        "                print(\"Set dry_run=False to perform the deletion.\")\n",
        "                print(\"IMPORTANT: Make sure you own all deployments before deleting.\")\n",
        "    else:\n",
        "        print(\"No deployments found.\")\n",
        "\n",
        "\n",
        "dry_run_delete_step = True  # @param {type:\"boolean\"}\n",
        "cleanup_deployments(dry_run_delete_step)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DLNQyaOjLfAY"
      },
      "source": [
        "Now, we'll deploy the multi-agent application to Vertex AI Agent Engine. This process might take a few minutes as it involves setting up the necessary infrastructure for remote execution.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4PK1o7dUXOdY"
      },
      "outputs": [],
      "source": [
        "# @title Define GCS bucket directory to stage the deployment artifacts\n",
        "artifacts_gcs_dirname = \"ag2_research_app\"  # @param {'type': 'string'}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3KXt3Kg7q08K"
      },
      "outputs": [],
      "source": [
        "# @title Deploy to Agent Engine - Please Wait\n",
        "agent_engine = AgentEngine.create(\n",
        "    ResearchApp(project_id=PROJECT_ID, location=LOCATION),\n",
        "    display_name=\"AG2 Research Application\",\n",
        "    gcs_dir_name=artifacts_gcs_dirname,\n",
        "    description=\"AG2 example deployed on Vertex AI Agent Engine\",\n",
        "    requirements=[\n",
        "        \"ag2[gemini]==0.8.3\",\n",
        "        \"google-cloud-aiplatform[agent_engines,ag2]==1.83.0\",\n",
        "        \"google-generativeai==0.8.4\",\n",
        "        \"dask[dataframe]==2024.12.1\",\n",
        "        \"cloudpickle==3\",\n",
        "    ],\n",
        "    extra_packages=[],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "18I5QCPkagp2"
      },
      "source": [
        "### Monitoring Deployment Progress with Google Cloud Logs Explorer\n",
        "\n",
        "You can track the progress of your deployment by examining the logs in the Google Cloud Logs Explorer - [Click Here!](https://console.cloud.google.com/logs/query;query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22). A successfully progressing deployment will typically exhibit log entries indicating the initialization of agents, the exchange of messages between agents, and the completion of various stages in the research process.\n",
        "\n",
        "**Here's what to look for in the logs:**\n",
        "\n",
        "* **Agent initialization messages:** Logs confirming that each agent (Researcher, Literature Reviewer, etc.) has been successfully initialized.\n",
        "* **Conversation logs:** Entries showing the back-and-forth communication between agents as they collaborate on the research task.\n",
        "* **Progress indicators:** Messages signaling the completion of key steps, such as data analysis, fact-checking, and literature review generation.\n",
        "* **Absence of error messages:** A lack of error or warning messages suggests that the deployment is proceeding without issues.\n",
        "\n",
        "\n",
        "Monitoring deployment progress through the deployment logs [here](https://console.cloud.google.com/logs/query;query=resource.type%3D%22aiplatform.googleapis.com%2FReasoningEngine%22) provides valuable insights into the health and functionality of your multi-agent application."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDmacPdS1ONB"
      },
      "source": [
        "### Testing the Deployed Multi-Agent Application\n",
        "\n",
        "Now that the multi-agent application is deployed to Vertex AI Agent Engines, we can test its functionality. We expect the interaction flow and results to be consistent with the local testing performed earlier. This ensures that the deployed application behaves as intended and provides the same user experience."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e8q38AWUwjQI"
      },
      "outputs": [],
      "source": [
        "result = agent_engine.query(topic=research_topic)\n",
        "display(Markdown(result))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zOt1dlPH5a5P"
      },
      "source": [
        "### Writing Report To Disk (Optional)\n",
        "\n",
        "The comprehensive literature review generated by the multi-agent conversation can be saved to a file. This allows for easy access, review, and potential modifications by users or other stakeholders.\n",
        "\n",
        "**Here's how you can persist the report to disk:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zz15hGQRwj4e"
      },
      "outputs": [],
      "source": [
        "with open(\"report.md\", \"w\") as f:\n",
        "    f.write(str(result))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "83m7sycmZ9GI"
      },
      "source": [
        "## Step 4: Clean up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WuqHrWIXakMa"
      },
      "source": [
        "To avoid incurring unnecessary costs, we can now safely remove the deployed Agent Engines."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T2Vf4H2ZqtSw"
      },
      "outputs": [],
      "source": [
        "dry_run_delete_step = False  # @param {type:\"boolean\"}\n",
        "cleanup_deployments(dry_run_delete_step)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A1XZlgAKan8I"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "This notebook demonstrated how to build and deploy a multi-agent AI system on Vertex AI using AG2 (formerly AutoGen) and its Agent Engines. We showcased a collaborative workflow where specialized agents worked together to perform complex research and analysis.\n",
        "\n",
        "### Agent Roles and Responsibilities\n",
        "\n",
        "The system comprises the following agents, each with a distinct role:\n",
        "\n",
        "*   **Researcher:** Conducts in-depth research on the given topic, identifying key trends, publications, experts, and relevant data sources. The Researcher provides accurate and well-cited information as the foundation for subsequent analysis.\n",
        "\n",
        "*   **Literature Reviewer:** Synthesizes the Researcher's findings into a coherent and structured literature review. This agent identifies key themes, connections, and gaps in the existing research to provide a comprehensive overview of the current state of knowledge.\n",
        "\n",
        "*   **Fact Checker:** Verifies the accuracy and validity of information presented by other agents, particularly the Researcher and Literature Reviewer. The Fact Checker meticulously checks sources, identifies potential biases, and ensures consistency with established knowledge.\n",
        "\n",
        "*   **Data Analyst:** Analyzes data related to the research topic, extracting meaningful insights, identifying trends, patterns, and correlations. The Data Analyst provides quantitative support for the research with clear and concise presentations of findings.\n",
        "\n",
        "*   **User Proxy:** Represents the user, initiating the conversation with the initial prompt or task. This agent can also execute code to support other agents, acting as an interface between the user and the multi-agent system. It also uses Gemini to reflect and create a final report using a custom prompt.\n",
        "\n",
        "The `GroupChatManager` orchestrates the interaction between these agents. It selects an agent to contribute, broadcasts the message to all other agents, and repeats this process until the conversation concludes. The default \"auto\" selection strategy utilizes the manager's LLM to intelligently choose the next speaker based on the conversation's context. This collaborative approach enables more comprehensive and robust analyses than a single agent could achieve.\n",
        "\n",
        "### Key Takeaways\n",
        "\n",
        "*   **Collaborative Multi-Agent Workflow:** Demonstrated a workflow where distinct agents (Researcher, Data Analyst, Literature Reviewer, Fact Checker, and User Proxy) collaborate to conduct research, analyze data, synthesize findings, and ensure accuracy.\n",
        "*   **Vertex AI Deployment:** Showcased the deployment and execution of this multi-agent system on Vertex AI Agent Engine, providing scalable and efficient cloud-based execution.\n",
        "*   **Simplified Agent Development with AG2:** Highlighted AG2's framework for building and managing AI agents, streamlining development and facilitating complex agent interactions.\n",
        "\n",
        "This example provides a foundation for building and deploying custom multi-agent applications on Vertex AI, leveraging the power of AG2 for collaborative AI solutions.\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tutorial_ag2_on_agent_engine.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
