gemini/agents/research-multi-agents/intro_research_multi_agents_gemini_2_0.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jJZw3h2myqls"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGrXr6X4yXG5"
      },
      "source": [
        "# Building a Research Multi Agent System - a Design Pattern Overview with Gemini 2.0\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/agents/research-multi-agents/intro_research_multi_agents_gemini_2_0.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%2Fagents%2Fresearch-multi-agents%2Fintro_research_multi_agents_gemini_2_0.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/agents/research-multi-agents/intro_research_multi_agents_gemini_2_0.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/agents/research-multi-agents/intro_research_multi_agents_gemini_2_0.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KJWqTM-CS0qC"
      },
      "source": [
        "<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/agents/intro_research_multi_agents_gemini_2_0.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/agents/intro_research_multi_agents_gemini_2_0.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/agents/intro_research_multi_agents_gemini_2_0.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/53/X_logo_2023_original.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/agents/intro_research_multi_agents_gemini_2_0.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/agents/intro_research_multi_agents_gemini_2_0.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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EnoVKOgny2ZM"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author(s) |  [Lavi Nigam](https://github.com/lavinigam-gcp)|"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CSt0qUR2Sg61"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>\n",
        "⚠️ Gemini 2.0 Flash (Model ID: gemini-2.0-flash-exp) and the Google Gen AI SDK are currently experimental and output can vary ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IDDZhYrClJQK"
      },
      "source": [
        "## Overview\n",
        "\n",
        "In today's rapidly evolving technology landscape, businesses frequently need to conduct comprehensive research and analysis that spans multiple data sources, requires complex reasoning, and demands clear actionable insights. Whether it's market research, competitive analysis, urban planning, or scientific research, the challenges remain similar: how to efficiently gather, process, and synthesize information while ensuring accuracy and scalability.\n",
        "\n",
        "In this notebook, as a developer, you'll discover how to create intelligent agents and multi-agent systems using Vertex AI Gemini 2.0.\n",
        "\n",
        "\n",
        "### Learning Through Implementation\n",
        "\n",
        "Rather than using existing frameworks, we'll build our multi-agent system from scratch. This approach offers several benefits:\n",
        "\n",
        "1. **Core Understanding**: Building from the ground up helps you understand the fundamental principles of multi-agent systems\n",
        "2. **Design Pattern Mastery**: Learn reusable patterns that work across different domains and technologies\n",
        "3. **Custom Control**: Gain the ability to fine-tune every aspect of your system\n",
        "4. **Debugging Confidence**: Understanding the internals makes troubleshooting much more straightforward\n",
        "\n",
        "While there are excellent open-source frameworks available for building multi-agent systems, such as [AutoGen](https://github.com/microsoft/autogen), [CrewAI](https://github.com/crewAIInc/crewAI), [PydanticAI](https://github.com/pydantic/pydantic-ai), and [LangGraph](https://github.com/langchain-ai/langgraph), we believe that a from-scratch approach in this notebook will provide a deeper understanding of the underlying concepts and mechanics.\n",
        "\n",
        "The open-source frameworks offers many valuable features like conditional routing, annotated global state, checkpointing, and more.\n",
        "\n",
        "Once you've grasped the fundamentals from this notebook, exploring these frameworks can unlock even more advanced capabilities and streamline your development process.\n",
        "\n",
        "\n",
        "### Key Technical Components\n",
        "\n",
        "Our implementation showcases essential Vertex AI ***Gemini 2.0*** capabilities:\n",
        "\n",
        "1. **Function Calling**: Structure agent behaviors and interactions\n",
        "2. **Structured Output**: Generate consistent, validatable data\n",
        "3. **Async Operations**: Handle parallel agent tasks efficiently\n",
        "4. **Google Search Integration**: Ground agent reasoning in real-world data\n",
        "\n",
        "\n",
        "### To get started, let's explore some key questions:\n",
        "\n",
        "* What exactly is an agent, and how does it differ from a simple LLM call?\n",
        "* How can agents use tools to achieve their goals?\n",
        "* And what possibilities emerge when multiple agents work together in a multi-agent system?\n",
        "\n",
        "\n",
        "#### **LLM Execution (The Foundation)**\n",
        "\n",
        "Think of an LLM as a powerful prediction engine. Given some input text (a prompt), it predicts what comes next, generating text, translating languages, writing different kinds of creative content, and answering your questions in an informative way. However, on its own, it simply reacts to your input and provides an output. It doesn't have a sense of purpose or the ability to act independently.\n",
        "\n",
        "**Example:** An LLM is like a super smart travel guidebook. You ask it \"What are some popular attractions in Paris?\" and it gives you a list. It provides information but doesn't actually do anything.\n",
        "\n",
        "![title](https://storage.googleapis.com/github-repo/generative-ai/gemini2/use-cases/research_multi_agent_ev/img/simple-llm-flow.png)\n",
        "\n",
        "#### **Agent (LLM with a Purpose)**\n",
        "\n",
        "Now, imagine giving that prediction engine some goals and the ability to act on them. This is essentially what an agent is. It's an LLM wrapped with extra code that allows it to:\n",
        "\n",
        "* **Understand the goal:** \"Book a flight to London.\"\n",
        "* **Break it down into steps:** Search for flights, compare prices, choose a date, make a booking.\n",
        "* **Use tools to achieve those steps:** Access a flight booking API, a web browser, or even interact with a human.\n",
        "\n",
        "**Example:** An agent is like a personal travel assistant. You tell it \"Plan a trip to Paris for me next month.\" The agent uses its LLM \"brain\" to understand what that means, then uses tools like flight booking websites, hotel search engines, and even weather apps to create an itinerary.\n",
        "\n",
        "\n",
        "![title](https://storage.googleapis.com/github-repo/generative-ai/gemini2/use-cases/research_multi_agent_ev/img/agent-flow.png)\n",
        "\n",
        "#### **Multi-Agent (Teamwork Makes the Dream Work)**\n",
        "\n",
        "Now, imagine several of these specialized agents working together, each with its own skills and responsibilities. That's a multi-agent system. They can communicate, share information, and coordinate their actions to achieve a complex goal.\n",
        "\n",
        "**Example:** Now imagine a team of specialized travel agents working together. One agent books the flights, another finds the perfect hotel, a third arranges tours and activities. They communicate and coordinate to create an amazing Paris trip.\n",
        "\n",
        "![title](https://storage.googleapis.com/github-repo/generative-ai/gemini2/use-cases/research_multi_agent_ev/img/multi-agent-flow.png)\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "Now that you have learned the fundamentals, moving forward, you'll learn the core design patterns behind agents and multi-agent systems. We'll demonstrate its capabilities through a practical use case - Electric Vehicle (EV) infrastructure expansion analysis - while keeping the core architecture adaptable for any research-intensive application."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ItHSSWRBG_D"
      },
      "source": [
        "## Objective\n",
        "\n",
        "This notebook will guide you through building a research-focused multi-agent system. Here's what you'll learn:\n",
        "\n",
        "* **A design pattern for creating these systems:** We'll introduce a reusable structure for building multi-agent systems geared towards research tasks.\n",
        "* **A practical example: EV Research Agent:**  See how we applied the design pattern to create an agent specializing in Electronic Vehicle research. This agent can answer complex queries like \"EV Charging Station Expansion in [City Name]\" by planning, researching, and generating a comprehensive report.\n",
        "* **Component integration and orchestration:** Understand how individual components within the agent work together seamlessly to produce the final output."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DKeB9J-EPFeQ"
      },
      "source": [
        "## Our Use Case: EV Infrastructure Analysis\n",
        "\n",
        "To demonstrate the power and flexibility of our Research Multi-Agent system, we'll tackle a real-world challenge: analyzing optimal locations for expanding Electric Vehicle (EV) charging infrastructure in cities across the United States.\n",
        "\n",
        "### The Challenge\n",
        "\n",
        "Urban planners and EV infrastructure companies face complex decisions when expanding charging networks:\n",
        "- Understanding population density and movement patterns\n",
        "- Analyzing existing charging infrastructure\n",
        "- Evaluating proximity to major highways and transit routes\n",
        "- Considering local demographics and economic factors\n",
        "- Assessing grid capacity and infrastructure readiness\n",
        "\n",
        "### Our Solution\n",
        "\n",
        "We'll build a research system that:\n",
        "1. Accepts queries about specific cities or regions\n",
        "2. Gathers data from multiple sources (OpenStreetMap, NREL API)\n",
        "3. Analyzes infrastructure patterns and gaps\n",
        "4. Generates actionable insights with citations\n",
        "5. Visualizes findings for better decision-making\n",
        "\n",
        "Simply, A team of research agents armed with data and search engines, technical know how, coordinated with a common goal across specialized skillsets and tasks.\n",
        "\n",
        "\n",
        "While we focus on EV infrastructure, the patterns and approaches we develop can be applied to any research-intensive domain requiring similar data gathering, analysis, and insight generation capabilities."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fAlzMGTdGhO-"
      },
      "source": [
        "## Gemini 2.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PkgSZ0wTzFU7"
      },
      "source": [
        "## Overview\n",
        "\n",
        "[Gemini 2.0 Flash](https://cloud.google.com/vertex-ai/generative-ai/docs/gemini-v2) is a new multimodal generative ai model from the Gemini family developed by [Google DeepMind](https://deepmind.google/). It now available as an experimental preview release through the Gemini API in Vertex AI and Vertex AI Studio. The model introduces new features and enhanced core capabilities:\n",
        "\n",
        "- Multimodal Live API: This new API helps you create real-time vision and audio streaming applications with tool use.\n",
        "- Speed and performance: Gemini 2.0 Flash is the fastest model in the industry, with a 3x improvement in time to first token (TTFT) over 1.5 Flash.\n",
        "- Quality: The model maintains quality comparable to larger models like Gemini 1.5 Pro and GPT-4o.\n",
        "- Improved agentic experiences: Gemini 2.0 delivers improvements to multimodal understanding, coding, complex instruction following, and function calling.\n",
        "- New Modalities: Gemini 2.0 introduces native image generation and controllable text-to-speech capabilities, enabling image editing, localized artwork creation, and expressive storytelling.\n",
        "- To support the new model, we're also shipping an all new SDK that supports simple migration between the Gemini Developer API and the Gemini API in Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "90JzDyyRzRRU"
      },
      "source": [
        "## Getting Started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qE48lDlSzf81"
      },
      "source": [
        "### Install Google Gen AI SDK for Python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "db8O7nh0zw_B"
      },
      "outputs": [],
      "source": [
        "# Downloading Google Gen AI SDK (experimental)\n",
        "%pip install google-genai\n",
        "\n",
        "# Libraries required for saving markdowns as external files.\n",
        "! apt install pandoc\n",
        "! apt install libreoffice"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2jCwQQxO0WVx"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sDXGN26_0Y0R"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IGcO4hXDzzuH"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you are running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rbm_CqxKz1b6"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ES6LwuBr0GSD"
      },
      "source": [
        "### Connect to a generative AI API service\n",
        "\n",
        "Google Gen AI APIs and models including Gemini are available in the following two API services:\n",
        "\n",
        "- **[Google AI for Developers](https://ai.google.dev/gemini-api/docs)**: Experiment, prototype, and deploy small projects.\n",
        "- **[Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/overview)**: Build enterprise-ready projects on Google Cloud.\n",
        "\n",
        "The Google Gen AI SDK provides a unified interface to these two API services.\n",
        "\n",
        "This notebook shows how to use the Google Gen AI SDK with the Gemini API in Vertex AI."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pMegXbM90JEk"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4mrov4hC0OZ-"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mp0umgC00TMZ"
      },
      "source": [
        "### Import libraries\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5FzK2TuA0SYe"
      },
      "outputs": [],
      "source": [
        "from google import genai\n",
        "from rich import print as rich_print"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C2iaXGH21j_U"
      },
      "source": [
        "### Create Gen AI Client"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "maA6ZXozxphR"
      },
      "outputs": [],
      "source": [
        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B70wFwV61uiK"
      },
      "source": [
        "### Load the Gemini 2.0 Flash model\n",
        "\n",
        "To learn more about all [Gemini models on Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#gemini-models)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9qipKyrW1vG9"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"gemini-2.0-flash-exp\"  # @param {type: \"string\"}\n",
        "MODEL_ID_Flash = \"gemini-1.5-flash-002\"  # For control generation for grounding with google search as a Tool"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oMdna9zsJKx7"
      },
      "source": [
        "To access comprehensive EV infrastructure data, you'll need an API key from the National Renewable Energy Laboratory (NREL). This key allows you to retrieve detailed information about EV charging stations, which is crucial for the `DataGatherAgent` to function correctly.\n",
        "\n",
        "**Here's how to get your NREL API key:**\n",
        "\n",
        "1. **Sign up:** Visit the [NREL Developer Network signup page](https://developer.nrel.gov/signup/).\n",
        "2. **Email Confirmation:** You'll receive an email with your API key.\n",
        "3. **Wait Time:** It might take some time to receive the email, so please be patient.\n",
        "4. **Check Spam:** Make sure to check your spam or junk folder if you don't see the email in your inbox.\n",
        "\n",
        "**Enter your API key in the following code cell:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HGKLVzF4JLJS"
      },
      "outputs": [],
      "source": [
        "NREL_API_KEY = \"[your-nrel-api-key]\"  # @param {type: \"string\", placeholder: \"[your-nrel-api-key]\", isTemplate: true}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nscpdicqtHa8"
      },
      "source": [
        "### Download utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IyWeX1ZqbVMZ"
      },
      "source": [
        "To streamline the process and keep our focus on the design, utility, and output of the multi-agent system, we've placed the core code for the `ev_agent` in an external location. This includes both the `agent_handler` and `api_handler`, which contain the main logic. However, we're now downloading it to our current environment to ensure we can import the necessary functions for our analysis:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P4cULOECtKeU"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/GoogleCloudPlatform/generative-ai.git \\\n",
        "  && cp -r generative-ai/gemini/agents/research-multi-agents/ev_agent ./ \\\n",
        "  && rm -rf generative-ai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ul4QvAKDbc6t"
      },
      "source": [
        "This makes the code, including all the agents and API handlers, readily available for use. You can always explore the downloaded code and make changes as you see fit. This approach allows us to keep the notebook cleaner and focused on the higher-level aspects of the system while still providing access to the underlying implementation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UYMXmaWny0JH"
      },
      "outputs": [],
      "source": [
        "# @title Saving Report (DOCX/PDF) Helper Functions\n",
        "\n",
        "import os\n",
        "import subprocess\n",
        "\n",
        "\n",
        "def convert_markdown(markdown_text, output_path, filename, file_type):\n",
        "    \"\"\"\n",
        "    Converts markdown text to DOCX or PDF using pandoc.\n",
        "\n",
        "    Args:\n",
        "        markdown_text: The markdown text to convert.\n",
        "        output_path: The directory where the output file should be saved.\n",
        "        filename: The name of the output file (without extension).\n",
        "        file_type: The desired output file type ('docx' or 'pdf').\n",
        "\n",
        "    Raises:\n",
        "        ValueError: If an invalid file type is specified.\n",
        "        FileNotFoundError: If pandoc is not found in the system's PATH.\n",
        "        subprocess.CalledProcessError: If the pandoc command fails.\n",
        "        OSError: If there is an error during file operations.\n",
        "    \"\"\"\n",
        "    os.makedirs(output_path, exist_ok=True)\n",
        "\n",
        "    if file_type not in [\"docx\", \"pdf\"]:\n",
        "        raise ValueError(\"Invalid file type specified. Must be 'docx' or 'pdf'.\")\n",
        "\n",
        "    docx_filepath = os.path.join(output_path, f\"{filename}.docx\")\n",
        "\n",
        "    try:\n",
        "        # Check if pandoc is available\n",
        "        subprocess.run([\"pandoc\", \"--version\"], capture_output=True, check=True)\n",
        "\n",
        "        # Convert Markdown to DOCX\n",
        "        subprocess.run(\n",
        "            [\"pandoc\", \"-f\", \"markdown\", \"-t\", \"docx\", \"-o\", docx_filepath],\n",
        "            input=markdown_text,\n",
        "            encoding=\"utf-8\",\n",
        "            check=True,\n",
        "        )\n",
        "        # print(f\"DOCX file saved to: {docx_filepath}\")\n",
        "\n",
        "        if file_type == \"pdf\":\n",
        "            pdf_filepath = os.path.join(output_path, f\"{filename}.pdf\")\n",
        "            # Convert DOCX to PDF (using libreoffice on Colab)\n",
        "            subprocess.run(\n",
        "                [\n",
        "                    \"libreoffice\",\n",
        "                    \"--headless\",\n",
        "                    \"--convert-to\",\n",
        "                    \"pdf\",\n",
        "                    \"--outdir\",\n",
        "                    output_path,\n",
        "                    docx_filepath,\n",
        "                ],\n",
        "                check=True,\n",
        "            )\n",
        "            print(f\"PDF file saved to: {pdf_filepath}\")\n",
        "\n",
        "            # Delete the temporary DOCX file\n",
        "            os.remove(docx_filepath)\n",
        "            print(f\"Temporary DOCX file deleted: {docx_filepath}\")\n",
        "\n",
        "    except FileNotFoundError:\n",
        "        raise FileNotFoundError(\n",
        "            \"pandoc not found. Please ensure it is installed and in your system's PATH.\"\n",
        "        )\n",
        "    except subprocess.CalledProcessError as e:\n",
        "        raise subprocess.CalledProcessError(\n",
        "            e.returncode, e.cmd, output=e.output, stderr=e.stderr\n",
        "        )\n",
        "    except OSError as e:\n",
        "        raise OSError(f\"Error during file operations: {e}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DUHn-HALvWuB"
      },
      "source": [
        "# Gemini-Powered EV Research: A Multi-Agent Approach\n",
        "\n",
        "This section outlines a powerful multi-agent system designed for in-depth research on Electric Vehicle (EV) charging infrastructure in US cities. Built entirely using Gemini 2.0, this system showcases a streamlined approach to complex research tasks.\n",
        "\n",
        "**Core Idea:** We've assembled a team of specialized AI agents, each using Gemini 2.0, to automate and enhance the research process. This approach leverages Gemini's strengths in:\n",
        "\n",
        "*   **Function Calling:** Enables agents to trigger specific actions and tools facilitating seamless interaction.\n",
        "*   **Structured Generations:** Ensures consistent, predictable output from each agent, simplifying inter-agent communication.\n",
        "*   **Async Model Calling:** Allows agents to work concurrently, significantly speeding up research.\n",
        "*   **Google Search Grounding:** Keeps the research grounded in real-world data and up-to-date information.\n",
        "\n",
        "## System Architecture\n",
        "\n",
        "At the heart of our system lies a clear, modular architecture, visualized below:\n",
        "\n",
        "![research-multi-agent-desing-pattern](https://storage.googleapis.com/github-repo/generative-ai/gemini2/use-cases/research_multi_agent_ev/img/multi-agent-design-pattern.png)\n",
        "**Agent Breakdown:**\n",
        "\n",
        "The diagram illustrates the core components of our system:\n",
        "\n",
        "*   **User (Pink):** Initiates the research process by submitting a query.\n",
        "*   **ExecutionAgent (Pink):** The central orchestrator, managing the workflow, handling communication between agents, and ensuring smooth execution. It also handles error recovery, such as retries and alternative execution paths, to maintain system robustness.\n",
        "*   **Core Research Agents (Green):**\n",
        "    *   **PlanningAgent:**  The strategist, converting the user's query into a detailed, step-by-step research plan.\n",
        "    *   **QueryAnalysisAgent:** The interpreter, determining the specific data required and the desired output format (e.g., raw data, report, visualization).\n",
        "    *   **DataGatherAgent:** The collector, responsible for fetching data from external APIs. It leverages Gemini's search grounding to ensure data accuracy and relevance. This agent is designed to be adaptable to various data sources.\n",
        "    *   **ReportAgent:** The writer, transforming raw data into a comprehensive, well-structured report. It can incorporate search-based grounding for validation and supports multiple output formats.\n",
        "    *   **VisualizeAgent:** The illustrator, creating clear and insightful visualizations (charts, graphs) to represent the findings. It adapts its output based on data types and user requirements.\n",
        "*   **Research Output (Pink):** The final, comprehensive research product delivered to the user.\n",
        "*   **External Systems (Blue):**\n",
        "    *   **External APIs:** Data sources for the `DataGatherAgent`.\n",
        "    *   **Visualization Tools:** Libraries used by the `VisualizeAgent`.\n",
        "    *   **Document Tools:** Resources utilized by the `ReportAgent` for formatting and presentation.\n",
        "\n",
        "\n",
        "**Benefits of the Gemini-Powered Approach:**\n",
        "\n",
        "*   **Simplified Development:** Build the entire system using a single, powerful API – Gemini.\n",
        "*   **Native Functionality:** Leverage Gemini's built-in features for seamless agent interaction and consistent output.\n",
        "*   **Enhanced Performance:** Async model calling enables parallel processing, accelerating the research process.\n",
        "*   **Real-World Relevance:** Google Search grounding ensures your research is always based on the latest information.\n",
        "*   **Scalability and Flexibility:**\n",
        "    *   Easily add new agents for specialized tasks (e.g., sentiment analysis of EV adoption).\n",
        "    *   Modify existing agents to adapt to new data sources or research requirements.\n",
        "    *   The modular design allows independent scaling of different system components."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W9mEtr_yw0vZ"
      },
      "source": [
        "## Exploring the EV Agent in Action\n",
        "\n",
        "Now that you've seen the architecture, let's dive into the practical side and see how our EV Research Agent works. We'll explore two ways to interact with it:\n",
        "\n",
        "**1. The \"Black Box\" Experience: Witnessing the Magic**\n",
        "\n",
        "Imagine the entire multi-agent system as a single, powerful unit – the `EVAgent`. In this section, we'll treat it as a \"black box.\" You'll simply send it a research query, and watch as it works behind the scenes, delivering a comprehensive report in about 1-2 minutes.\n",
        "\n",
        "We'll try two exciting examples:\n",
        "\n",
        "*   **Example 1: Basic Report Generation:**  See how the agent generates a structured report with predefined sections based on your query.\n",
        "*   **Example 2: Google Search Enhanced Report:** Observe how the agent leverages Google Search to enrich the report with citations, deeper insights, and up-to-the-minute information.\n",
        "\n",
        "**2. Deconstructing the Process: A Step-by-Step Journey**\n",
        "\n",
        "Ready to peek under the hood? In this section, we'll dissect the agent's inner workings. You'll follow along as your query is processed through each stage of the research pipeline:\n",
        "\n",
        "*   **Planning:** Witness how the `PlanningAgent` crafts the initial research strategy. *We'll briefly touch upon the code behind this, highlighting the input it receives and the plan it outputs, along with the data models that structure this communication.*\n",
        "*   **Reasoning:** See how the `QueryAnalysisAgent` determines the necessary data and output format. *Again, we'll peek at the underlying code to understand its input, output, and the data models involved.*\n",
        "*   **Tool Selection:** Observe how the `DataGatherAgent` chooses the right APIs and leverages Google Search. *We'll examine the code's role in this selection process, focusing on the data models that guide its choices.*\n",
        "*   **Coordination:** Understand how the `ExecutionAgent` orchestrates the entire process. *We will shed some light on the code that enables this coordination, emphasizing the data models as the communication backbone between agents.*\n",
        "*   **Decision-Making:** Learn how the agents make choices at each step, leading to the final output.\n",
        "\n",
        "You'll see firsthand how these individual steps, powered by their underlying logic and data models, contribute to the final, polished report and visualizations.\n",
        "\n",
        "**A Note on Code Structure:**\n",
        "\n",
        "To keep this exploration clear and focused, the detailed code for each agent is neatly organized in separate files. **We are choosing not to put code directly in this notebook as it will make it unnecessarily complex.** So when we go through step by step, think of each agent as a black box. We will, however briefly talk about the design pattern it follows, what the data model it uses behind the scene to produce an output. Once you understand that, you can easily refer to the code from scratch or use any open-source library to implement a similar agent. Think of them as behind-the-scenes appendices you can explore later to dive deep into the implementation details of each agent.\n",
        "\n",
        "**The primary goal here is to showcase the power of agent collaboration with Gemini 2.0.** You'll witness how our team of Gemini-powered agents works together seamlessly to fulfill your research requests, demonstrating the elegance and efficiency of this multi-agent approach.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Xh-cZwQFJpz"
      },
      "source": [
        "## EV Agent - The \"Black Box\" Experience:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "on2mGosD1WBp"
      },
      "source": [
        "The `ExecutionAgent` is the heart of our EV infrastructure analysis system. Think of it as the conductor of an orchestra, coordinating a team of specialized agents to perform a comprehensive analysis based on your query.\n",
        "\n",
        "**Before you start:**\n",
        "\n",
        "*   **What it does:**  The `ExecutionAgent` takes your query about EV infrastructure, develops a plan, gathers relevant data, generates reports, and creates insightful visualizations.\n",
        "*   **How it works:** It delegates tasks to other agents (like a planning agent, data gathering agent, etc.) and manages the overall workflow.\n",
        "*   **What you get:** You'll receive a structured output containing the analysis plan, gathered data, a detailed report (if requested), and visualizations (if applicable).\n",
        "*   **Customization:** You can control the level of detail (debug mode), whether to see intermediate outputs (stage\\_output), and the type of output you desire (e.g., raw data, report, text).\n",
        "\n",
        "Essentially, the `ExecutionAgent` simplifies the complex process of EV infrastructure analysis, providing you with a powerful tool to gain valuable insights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K5O-SIBr0g7p"
      },
      "outputs": [],
      "source": [
        "# Importing ExecutionAgent from our agent_handler\n",
        "\n",
        "\n",
        "from ev_agent.agent_handler.agent_01_ExecutionAgent import ExecutionAgent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EeXrJaqTIBJv"
      },
      "outputs": [],
      "source": [
        "# Create the agent\n",
        "\n",
        "agent = ExecutionAgent.create(\n",
        "    client=client,\n",
        "    model_name=MODEL_ID_Flash,  # Gemini 2.0 Flash\n",
        "    api_key=NREL_API_KEY,\n",
        "    debug=False,\n",
        "    stage_output=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YgUWghr5IG5T"
      },
      "source": [
        "###  Basic Report Generation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pRvAIJmD2p76"
      },
      "source": [
        "In this case, we're treating the `ExecutionAgent` as a **\"black box\"**. We provide the input query (\"I want to understand the EV charging situation in Austin.\") and it will eventually deliver the final report without revealing the inner workings."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lgIT_0Wc2uPK"
      },
      "source": [
        "Since we set `debug=False` and `stage_output=False` earlier, the agent is giving us some playful warnings. It's essentially saying, \"Hey, you've turned off all the visibility into the process, so you'll only see the final result! But, just so you know, there are four agents working hard behind the scenes\"."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QI6Lfz6QIek4"
      },
      "outputs": [],
      "source": [
        "# Execute the analysis\n",
        "results = await agent.execute(\n",
        "    \"I want to understand the EV charging situation in Austin.\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a4I0IohK2wPn"
      },
      "source": [
        "You'll notice a humorous warning: `*Deciphering your cryptic commands! It's like translating ancient hieroglyphs, but with more emojis.*` This is a subtle hint that the **QueryAnalysisAgent** is currently at work, interpreting your input query. If you ever want to peek behind the curtain, simply set `debug=True` or `stage_output=True` when creating the agent. But for now, we're embracing the black box experience and eagerly awaiting the final, comprehensive report."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fMWrJ7AyD5FW"
      },
      "source": [
        "---\n",
        "If you want to save the generated report for later use or sharing, you can easily convert it to PDF or DOCX format. Here's how:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q3n4cZJj5IjH"
      },
      "outputs": [],
      "source": [
        "# # You can save the report as PDF or DOCX\n",
        "\n",
        "markdown_text = (\n",
        "    results[\"report\"][\"full_text\"] + \"\\n\\n\\n\" + results[\"report\"][\"citations\"]\n",
        ")\n",
        "\n",
        "convert_markdown(\n",
        "    markdown_text,\n",
        "    output_path=\"/content/generated_report\",\n",
        "    filename=\"austin_normal\",\n",
        "    file_type=\"pdf\",  # or \"docx\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZXk8CAHE5pRE"
      },
      "source": [
        "This will generate a nicely formatted report file in your chosen location, ready to be viewed or shared. You can see an example of a pre-generated report here: [Austin Report with Sections](https://storage.googleapis.com/github-repo/generative-ai/gemini2/use-cases/research_multi_agent_ev/sample_reports/austin_normal.pdf)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7vb7X_RUFIPi"
      },
      "source": [
        "The `results` object is a dictionary containing all the data generated from the analysis, including the `plan`, `query_analysis`, `data`, and the final `report` (with `citations`, `full_text`, and `sections`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UFH-oLQDKAxP"
      },
      "outputs": [],
      "source": [
        "rich_print(\n",
        "    \"The result object contains all these internal data points with the reports: \",\n",
        "    list(results.keys()),\n",
        ")\n",
        "rich_print(\n",
        "    \"The Report contains the citations, full text of the report and individual sections: \",\n",
        "    list(results[\"report\"].keys()),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R8JZGlBcGBKw"
      },
      "source": [
        "We've saved the full report above, but for now, let's just look at one section to see how they're structured. This demonstrates the organized way we store information within the report.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Osqut3cyLFkx"
      },
      "outputs": [],
      "source": [
        "for section_name, section_text in results[\"report\"][\"sections\"].items():\n",
        "    if section_name == \"Infrastructure Overview\":\n",
        "        print(section_name)\n",
        "        rich_print(section_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JrGgCle_GsSr"
      },
      "source": [
        "Let's focus on the data structure of each report section, which is crucial for developers to understand, especially in the context of our multi-agent system.\n",
        "\n",
        "As shown in the output, each section, like \"Infrastructure Overview,\" is represented as a `Section` object. This object neatly encapsulates:\n",
        "\n",
        "*   **`title`:** The title of the section (e.g., \"Infrastructure Overview\").\n",
        "*   **`content`:** The main text of the section, generated by Gemini, providing a detailed analysis. It's important to note that this content is dynamically created based on the data gathered by the `DataGatherAgent` and the insights generated by the language model.\n",
        "*   **`citations`:** A dictionary containing `CitationData` objects. Each citation provides a `number`, `value`, `data_path`, `raw_value`, and `context`, meticulously linking claims in the content to specific data points retrieved by our `DataGatherAgent` via API calls.\n",
        "*   **`key_findings`:** A list of key insights extracted from the section's content.\n",
        "*   **`enhanced_content`:** An optional field for additional data or analysis.\n",
        "\n",
        "\n",
        "In the normal \"Infrastructure Overview\" section, the numbers and facts presented are not manually entered; they are dynamically derived from our structured data model. This model is populated with real-world data fetched from various APIs by our dedicated `DataGatherAgent`. Let's see how this works with an example:\n",
        "\n",
        "**From the \"Infrastructure Overview\" section:**\n",
        "\n",
        "> \"Austin's total area encompasses 1679.20 sq km [1], with a significant portion dedicated to built areas (644.59 sq km) [1].\"\n",
        "\n",
        "The numbers \"1679.20\" and \"644.59\" are linked to **Citation 1**:\n",
        "\n",
        "```\n",
        "1: CitationData(\n",
        "    number=1,\n",
        "    value='1679.20 sq km total area, 644.59 sq km built area, 42224 service roads, 476 EV charging\n",
        "stations',\n",
        "    data_path='summary.area_metrics.total_area, summary.area_metrics.built_area,\n",
        "summary.roads.service_roads, summary.parking.ev_charging',\n",
        "    raw_value=\"{'total_area': {'value': '1679.20', 'path': 'summary.area_metrics.total_area_sqkm', 'unit':\n",
        "'sq km'}, 'built_area': {'value': '644.59', 'path': 'summary.area_metrics.built_area_sqkm', 'unit': 'sq km'},\n",
        "'service_roads': {'value': '42224', 'path': 'summary.roads.service_roads', 'unit': 'roads'}, 'ev_charging':\n",
        "{'value': '476', 'path': 'summary.parking.ev_charging', 'unit': 'stations'}}\",\n",
        "    context='Overall Austin metrics and existing EV charging station count'\n",
        "),\n",
        "```\n",
        "\n",
        "**Here's the breakdown:**\n",
        "\n",
        "1. **Data Source:** The `DataGatherAgent` makes API calls to sources like OpenStreetMap to gather data about Austin.\n",
        "2. **Structured Data Model:** This fetched data is stored in a structured format. For example, `summary.area_metrics.total_area` is a specific field in our data model that holds Austin's total area.\n",
        "3. **Citation Tracing:** Citation 1 clearly links the numbers in the text to their source in the data model. The `data_path` field shows where to find the data (e.g., `summary.area_metrics.total_area`), and the `raw_value` field reveals the exact value fetched from the API (\"1679.20\").\n",
        "4. **Dynamic Content Generation:** When the report is generated, the system automatically pulls the relevant data from the model, based on the `data_path` specified in the citation, and inserts it into the text.\n",
        "\n",
        "**Why is this important?**\n",
        "\n",
        "*   **Accuracy:** Our report is based on real data from trusted APIs, not on manual input, minimizing errors.\n",
        "*   **Traceability:** We can always trace the data back to its source, ensuring transparency and verifiability.\n",
        "*   **Automation:** The `DataGatherAgent` and our structured data model automate the data retrieval and integration process, making it efficient.\n",
        "*   **Consistency:** This structured approach ensures consistency across the report, as all agents use the same data model.\n",
        "\n",
        "In essence, the normal section demonstrates the power of our data-driven approach. The `DataGatherAgent`, our structured data model, and the `CitationData` system work together seamlessly to create a report grounded in accurate, traceable, and automatically updated information. This highlights the core strength of our multi-agent system: its ability to leverage structured data to produce reliable and insightful analysis.\n",
        "\n",
        "\n",
        "\n",
        "**Why is this data structure useful for developers and a multi-agent system?**\n",
        "\n",
        "This structured format promotes modularity, allowing developers to reuse sections and enabling different agents to collaborate seamlessly by contributing to specific parts of the report. The clear link between generated content and underlying data via CitationData ensures data integrity and transparency. Furthermore, the design is extensible, accommodating future growth and new types of analysis without disrupting the core structure, making it ideal for a multi-agent system."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w7CchYhMItoe"
      },
      "source": [
        "While we've focused on the report, you can also explore other parts of the `results` object. This provides a way to delve deeper into the agent's inner workings, but we'll break down each agent's role in more detail in the next section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "926d0bb80f59"
      },
      "outputs": [],
      "source": [
        "# You can print the whole text of the report:\n",
        "rich_print(results[\"report\"][\"full_text\"])\n",
        "\n",
        "# You can print the whole citations of the report:\n",
        "rich_print(results[\"report\"][\"citations\"])\n",
        "\n",
        "# You can also check the data it has used to generate the report\n",
        "rich_print(results[\"data\"])\n",
        "\n",
        "# If you want to see the whole plan of the agent that it executed\n",
        "rich_print(results[\"plan\"])\n",
        "\n",
        "# If you want to see the query analysis of the agent\n",
        "rich_print(results[\"query_analysis\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kiumiZy46oob"
      },
      "source": [
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8JIQYmPkIL1q"
      },
      "source": [
        "### Google Search Enhanced Report"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YKw5q6woJ8pt"
      },
      "source": [
        "Now, let's kick it up a notch! We're going to run the analysis again with `results_grounded_plot = await agent.execute(\"\"\"I want to understand the EV charging situation in Austin. I need a report and enhance the sections of report with google. Also add some plots\"\"\")`. This time, we've added two new twists to our request: grounding the report sections with Google Search results and adding data plots.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t-3L1Nz-1nPy"
      },
      "source": [
        "**Note on grounding with google search as a Tool with Gemini 2.0:**\n",
        "\n",
        "Currently, grounding with google search as a Tool on Gemini 2.0 does not support controlled generation. While you can still perform grounding with search, the output format and structure cannot be explicitly controlled at this time. Controlled generation is important for grounding as it allows us to specify the desired format and structure of the output, ensuring that the information retrieved from web search is integrated into the report in a consistent and organized manner. In the meantime, we are utilizing the Gemini 1.5 Flash model to perform grounding with controlled generation capabilities. You can explore examples of grounding with google search as a Tool Gemini 2.0 (without controlled generation) [here](link)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rboGC6pP1EBW"
      },
      "outputs": [],
      "source": [
        "# Create the agent\n",
        "\n",
        "agent = ExecutionAgent.create(\n",
        "    client=client,\n",
        "    model_name=MODEL_ID_Flash,  # Gemini 1.5 Flash\n",
        "    api_key=NREL_API_KEY,\n",
        "    debug=False,\n",
        "    stage_output=False,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bCZ23K5SIMMj"
      },
      "outputs": [],
      "source": [
        "# Execute the analysis\n",
        "results_grounded_plot = await agent.execute(\n",
        "    \"\"\"I want to understand the EV charging situation in Austin. I need a report and enhance the sections of report with google. Also add some plots\"\"\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jzz2vbC1J_Q8"
      },
      "source": [
        "Just like before, you'll see the familiar playful warnings since we're still running in a \"black box\" mode. However, now you'll also notice `DEBUG` messages indicating that sections are being enhanced with new citations, for example: `DEBUG: Enhanced Executive Summary with 17 new citations`. This is where the magic happens! The agent is now smartly integrating information from Google Search to bolster the report.\n",
        "\n",
        "What can you expect? Not only will the report be more comprehensive and grounded in a wider range of sources, but you'll also get to see insightful visualizations of the data. This is a significant step up from the previous run, showcasing the agent's ability to dynamically adapt to our requests and provide a richer, more visually engaging analysis. Get ready to be impressed by the power of combining AI, data analysis, and web search in a single, seamless process!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "90xiSV6G6Jts"
      },
      "outputs": [],
      "source": [
        "# Just like before, you can save this enhanced report as a PDF or DOCX using:\n",
        "\n",
        "convert_markdown(\n",
        "    markdown_text=results_grounded_plot[\"report\"][\"combined_report\"],\n",
        "    output_path=\"/content/generated_report\",\n",
        "    filename=\"austin_grounded\",\n",
        "    file_type=\"pdf\",  # or \"docx\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CVB0HBPP7Rgu"
      },
      "source": [
        "This will generate a file with the grounded sections. If you're eager to see the complete report right away, you can check out the pre-generated version here: [Austin Report - Sections Grounded with Search](https://storage.googleapis.com/github-repo/generative-ai/gemini2/use-cases/research_multi_agent_ev/sample_reports/austin_grounded.pdf)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H1yg0GWZN17j"
      },
      "source": [
        "You've seen the full, enhanced report – now let's take a closer look at how a single grounded section compares to the normal section we saw earlier. We'll examine the \"Infrastructure Overview\" section again:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4nCGfx-NMdel"
      },
      "outputs": [],
      "source": [
        "for section_name, section_text in results_grounded_plot[\"report\"][\"sections\"].items():\n",
        "    if section_name == \"Infrastructure Overview\":\n",
        "        print(section_name)\n",
        "        rich_print(section_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1jLTWQe0N3ZY"
      },
      "source": [
        "**Here's the \"aha\" moment:** Notice how the `content` of this section is now significantly richer and more detailed. It's not just stating facts from our initial data; it's weaving in insights and information gathered from the web through Google Search. This demonstrates the power of grounding our analysis in a broader context."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uVlnk6OxQyn1"
      },
      "source": [
        "Okay, let's break down how Google Search enhances the report by focusing on a specific example: **Citation 8**.\n",
        "\n",
        "In the grounded \"Infrastructure Overview\" section, we have:\n",
        "\n",
        "```\n",
        "        8: CitationData(\n",
        "            number=8,\n",
        "            value=\"Report on global EV infrastructure trends and best practices. | Context: Informs strategic\n",
        "recommendations for improving Austin's EV infrastructure. | URL: BloombergNEF\",\n",
        "            data_path='BloombergNEF',\n",
        "            raw_value='Report on global EV infrastructure trends and best practices.',\n",
        "            context=\"Informs strategic recommendations for improving Austin's EV infrastructure.\"\n",
        "        )\n",
        "```\n",
        "\n",
        "This citation points to a report from **BloombergNEF** on global EV infrastructure trends. Now, let's see how this reference, found through Google Search, contributes to the enhanced content:\n",
        "\n",
        "**Original Content (Before Search):**\n",
        "\n",
        "> \"The existing EV charging infrastructure, while growing, needs significant expansion to meet the rising demand for EVs. Currently, there are 78 total EV charging stations [2] across the city. This number is significantly lower than other major cities with similar populations.\"\n",
        "\n",
        "**Enhanced Content (After Search):**\n",
        "\n",
        "> \"The existing EV charging infrastructure, while growing, needs significant expansion to meet the rising demand for EVs. Currently, there are 78 total EV charging stations [2] across the city. This number is significantly lower than other major cities with similar populations. **A recent study by BloombergNEF [3] highlights the need for a much higher density of charging stations to support widespread EV adoption.**\"\n",
        "\n",
        "**Here's the impact:**\n",
        "\n",
        "1. **External Validation:** The original content stated that Austin's charging station count is low compared to similar cities. The enhanced content, using the BloombergNEF report found via Google Search, adds external validation to this claim. It's no longer just an observation based on our data; it's now supported by a reputable source on global EV trends.\n",
        "2. **Strategic Depth:** The BloombergNEF citation adds a layer of strategic depth. It's not just about the current number of stations; it connects to the broader concept of \"charging station density\" needed for \"widespread EV adoption\" – a key insight for planning Austin's EV future.\n",
        "3. **Credibility Boost:** Referencing a well-known organization like BloombergNEF significantly enhances the credibility of the report. It demonstrates that our analysis is informed by industry experts and best practices.\n",
        "\n",
        "**In essence, Google Search, through this specific citation, helped us transform a simple observation into a well-supported, strategically relevant insight.** It demonstrates how our system leverages web knowledge to enhance the report's quality, moving beyond the limitations of our initial data and providing a more nuanced and impactful analysis. This dynamic integration of external information is a key strength of our multi-agent approach."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CZLXKarZVcQC"
      },
      "source": [
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hUrIIjc6To8S"
      },
      "source": [
        "Now, let's visualize the raw data that underpins our analysis. The following code will generate plots directly from the data fetched by our `DataGatherAgent` from external APIs.\n",
        "\n",
        "You can also check the data it has used to generate the plots\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7cd5d15ec423"
      },
      "outputs": [],
      "source": [
        "rich_print(results_grounded_plot[\"data\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CuELS1YtVNVA"
      },
      "source": [
        "Let's explore the visualizations generated from the raw API data, which offer a deeper understanding of Austin's EV infrastructure and its urban context. The code uses the `create_comprehensive_city_analysis` function to produce a set of Plotly figures, each shedding light on different aspects of the city:\n",
        "\n",
        "**1. EV Infrastructure Overview Dashboard:**\n",
        "\n",
        "*   **Charging Station Types:** This bar chart breaks down the number of DC Fast, Level 2, and Level 1 charging stations. For Austin, it highlights the dominance of Level 2 chargers and the relative scarcity of DC Fast chargers. This is crucial for understanding the current charging landscape and identifying potential gaps, especially for users requiring faster charging options.\n",
        "*   **Connector Distribution:** This pie chart reveals the types of connectors available (e.g., CCS, CHAdeMO, Tesla). By examining this chart for Austin, you can assess the compatibility of the existing infrastructure with various EV models.\n",
        "*   **Network Distribution:** This bar chart displays the number of charging stations associated with different networks (e.g., ChargePoint, Tesla). For Austin, it might reveal a reliance on a particular network, which could influence decisions about network diversification and partnerships.\n",
        "*   **Access & Payment Methods:** This bar chart shows the percentage of stations offering various access and payment methods (e.g., credit card, mobile pay, 24/7 access). In Austin's case, it can indicate the ease of use and accessibility of the charging infrastructure for different users.\n",
        "\n",
        "**2. Transportation Infrastructure Analysis:**\n",
        "\n",
        "*   **Public Transport Facilities:** This section visualizes the number of bus stops, train stations, bus stations, and bike rental locations. For Austin, this data helps assess the integration of EV charging with existing public transportation, which is vital for planning intermodal hubs.\n",
        "*   **Road Network Distribution:** This shows the distribution of motorways, primary, secondary, and residential roads. Understanding Austin's road network density and types can inform decisions about optimal charging station placement along major thoroughfares.\n",
        "*   **Parking Facilities:** This section charts the number of surface parking lots, parking structures, street parking spaces, and designated EV charging spots. For Austin, it helps evaluate the availability of parking spaces that could potentially be equipped with EV charging.\n",
        "*   **EV vs. Traditional Infrastructure:** This compares the number of EV charging stations, fuel stations, car dealerships, and car repair shops. In Austin's context, it provides insights into the current balance between EV and traditional vehicle infrastructure, indicating the progress of EV adoption.\n",
        "\n",
        "**3. Urban Amenities and Services:**\n",
        "\n",
        "*   **Retail and Shopping:** This visualizes the distribution of shopping centers, supermarkets, department stores, and convenience stores. For Austin, it helps identify potential locations for charging stations near high-traffic retail areas.\n",
        "*   **Food and Entertainment:** This section charts restaurants, cafes, bars, and fast-food outlets. Understanding the density of these amenities in Austin can guide the placement of charging stations near popular destinations.\n",
        "*   **Emergency Services:** This displays the number of police stations, fire stations, hospitals, and clinics. For Austin, this information can be relevant for ensuring the resilience of the EV infrastructure and planning for emergency response related to EVs.\n",
        "*   **Public Amenities:** This visualizes the number of post offices, banks, ATMs, and public toilets. In Austin's context, it helps assess the availability of essential services near potential charging station locations.\n",
        "\n",
        "**4. Area Analysis:**\n",
        "\n",
        "*   **Area Distribution:** This pie chart shows the breakdown of Austin's total area into water, green, built, and other areas. It provides a quick overview of the city's land use, which can be a factor in determining suitable locations for charging infrastructure.\n",
        "\n",
        "**Ideally, these charts would be integrated into the report itself, providing a visual complement to the textual analysis.** However, even as standalone visualizations, they offer valuable insights for decision-making related to EV charging station expansion. For example, by examining the distribution of charging types, connector types, and network providers, along with the city's transportation infrastructure and urban amenities, stakeholders can identify strategic locations for new charging stations, optimize the mix of charging options, and ensure that the expansion aligns with the city's overall development and EV adoption trends. By correlating the density of public transportation, road networks, and parking facilities with the location of existing EV charging stations, planners can pinpoint areas where additional infrastructure is most needed. They can also consider factors such as proximity to retail centers, food and entertainment venues, and public amenities to enhance the user experience and maximize the utilization of charging stations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Q_MfimBMdbS"
      },
      "outputs": [],
      "source": [
        "print(\"\\n=== Single City Analysis ===\")\n",
        "for name, fig in results_grounded_plot[\"visualizations\"][0].items():\n",
        "    print(f\"\\nDisplaying: {name.replace('_', ' ').title()}\")\n",
        "    fig.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YE6bl8haTPL-"
      },
      "source": [
        "The `results` object is a dictionary containing all the data generated from the analysis, includes extra variables to add visualizations and `combined_report`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K_j2N-_5Kv_D"
      },
      "outputs": [],
      "source": [
        "rich_print(\n",
        "    \"The result object contains all these internal data points with the reports: \",\n",
        "    list(results_grounded_plot.keys()),\n",
        ")\n",
        "rich_print(\n",
        "    \"The Report contains the combined reports, citations, full text of the report and individual sections: \",\n",
        "    list(results_grounded_plot[\"report\"].keys()),\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bz-FBnhJTi-5"
      },
      "source": [
        "While we've focused on the report, you can also explore other parts of the `results` object. This provides a way to delve deeper into the agent's inner workings, but we'll break down each agent's role in more detail in the next section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1192d611f92b"
      },
      "outputs": [],
      "source": [
        "# You can print the whole report:\n",
        "rich_print(results_grounded_plot[\"report\"][\"combined_report\"])\n",
        "\n",
        "# You can print the whole text of the report:\n",
        "rich_print(results_grounded_plot[\"report\"][\"full_text\"])\n",
        "\n",
        "# You can print the whole citations of the report:\n",
        "rich_print(results_grounded_plot[\"report\"][\"citations\"])\n",
        "\n",
        "# You can also check the data it has used to generate the report\n",
        "rich_print(results_grounded_plot[\"data\"])\n",
        "\n",
        "# If you want to see the whole plan of the agent that it executed\n",
        "rich_print(results_grounded_plot[\"plan\"])\n",
        "\n",
        "# If you want to see the query analysis of the agent\n",
        "rich_print(results_grounded_plot[\"query_analysis\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GcW6BJTq8WuT"
      },
      "source": [
        "---\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jGdQeo_CRHh_"
      },
      "source": [
        "## Deconstructing the Process: A Step-by-Step Journey of Agents"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ruk0__600E3e"
      },
      "source": [
        "Before we delve into the inner workings of each agent, let's take a look at the overall flow of our multi-agent system. This sequence diagram provides a visual representation of how the agents interact and collaborate to process your query and generate the final output:\n",
        "\n",
        "![research-multi-agent-desing-pattern](https://storage.googleapis.com/github-repo/generative-ai/gemini2/use-cases/research_multi_agent_ev/img/ev_agent_simple.png)\n",
        "\n",
        "\n",
        "This sequence diagram serves as a visual roadmap for understanding the flow of our multi-agent system, and you can refer back to it as we explore each agent's inner workings. It illustrates how agents like the `ExecutionAgent`, `PlanningAgent`, `QueryAnalysisAgent`, `DataGatherAgent`, `ReportAgent`, and `VisualizeAgent` interact and collaborate to process your query, highlighting their roles, the flow of information, and key decision points. This diagram is crucial for grasping the big picture as we delve into the specifics of each agent, starting with the `PlanningAgent`, which initiates the analysis process based on your query."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hIQNRrF9KWqL"
      },
      "source": [
        "### Agent: PlanningAgent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BHYuwLIgDHyw"
      },
      "source": [
        "### Agent: PlanningAgent\n",
        "\n",
        "The `PlanningAgent` is the first active agent in our sequence, responsible for taking your initial query and crafting a strategic execution plan. As seen in the sequence diagram, the `ExecutionAgent` passes the user's query to the `PlanningAgent`, which then returns a structured plan. Let's break down its role:\n",
        "\n",
        "**Input:**\n",
        "\n",
        "*   **Query:** The user's raw query about EV infrastructure (e.g., \"Analyze EV charging stations in Austin\").\n",
        "*   **Client:** An instance of the generative AI model client (e.g., `gemini`).\n",
        "*   **Model Name:** The specific model to be used (e.g., \"gemini-pro\").\n",
        "*   **Debug:** A boolean flag to enable/disable debug mode.\n",
        "*   **API Key:** The API key for external services like NREL.\n",
        "\n",
        "**Output:**\n",
        "\n",
        "*   **ExecutionPlan:** A structured plan containing:\n",
        "    *   **Query:** The original user query.\n",
        "    *   **Timestamp:** When the plan was created.\n",
        "    *   **Validated Query:** Result of query validation, including validity, cities mentioned, missing elements, and suggestions for improvement.\n",
        "    *   **Enable Search:** A boolean flag indicating if enhanced search/grounding is required.\n",
        "    *   **Steps:** A list of `PlanStep` objects, each defining a step in the execution process with details like agent name, description, input/output formats, and status.\n",
        "    *   **Debug:**  A boolean flag indicating debug status.\n",
        "\n",
        "This section will explore five key aspects of the `PlanningAgent`: its setup, the creation of the `ExecutionPlan`, query validation and suggestions, handling of invalid queries, and a glimpse into its internal code structure."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qIj_pscoDmQU"
      },
      "source": [
        "Agent Code:\n",
        "```\n",
        "`/content/ev_agent/agent_handler/agent_02_PlanningAgent.py`\n",
        "```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qLxSAgN3AyyW"
      },
      "source": [
        "#### Setting up and Calling the agent\n",
        "\n",
        "First, we need to set up and call the `PlanningAgent`. Here's how we do it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6ZzI1mSbJ3bt"
      },
      "outputs": [],
      "source": [
        "from ev_agent.agent_handler.agent_02_PlanningAgent import *"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pFrQSpP5E3Kv"
      },
      "source": [
        "We start by importing the necessary `PlanningAgent` class. Then, we create an instance of the agent, providing the user's query, the client object, the model name, and setting `debug` to `False` for now. Finally, we call the `create_plan()` method to generate the execution plan. If `debug` is set to `False`, you might see a humorous warning about the complexity of plan creation, which is just a playful way to indicate that the agent is working behind the scenes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qONvCmwqJ3Yq"
      },
      "outputs": [],
      "source": [
        "agent = PlanningAgent(\n",
        "    query=\"I want to understand the EV charging situation in austin and proper vetted information and some plot\",\n",
        "    client=client,\n",
        "    model_name=MODEL_ID_Flash,\n",
        "    debug=False,\n",
        ")\n",
        "plan = agent.create_plan()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GQjRth99ATj0"
      },
      "source": [
        "#### ExecutionPlan"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7BOaeF9aGULR"
      },
      "source": [
        "Now, let's examine the `ExecutionPlan` generated by the `PlanningAgent`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VGN-ZitTJ3Vz"
      },
      "outputs": [],
      "source": [
        "rich_print(plan)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rLe_ptdcGaf1"
      },
      "source": [
        "The core of this plan lies in the `steps` list, which contains a sequence of `PlanStep` objects. Each `PlanStep` is defined by a structured data model, specifying:\n",
        "\n",
        "*   **`step_id`:** A unique identifier for the step.\n",
        "*   **`agent_name`:** The name of the agent responsible for this step (e.g., `QueryAnalysisAgent`, `DataGatherAgent`).\n",
        "*   **`description`:** A brief description of the step's purpose.\n",
        "*   **`input_requirements`:** The data required for this step (e.g., the output of a previous step).\n",
        "*   **`output_format`:** The format of the data produced by this step (e.g., a specific data model like `QueryEntity` or `DataGatherAgentOutput`).\n",
        "*   **`status`:** The current status of the step (e.g., `PENDING`, `COMPLETED`).\n",
        "*   **`error`:** Any error encountered during the step (initially `None`).\n",
        "*   **`skip_conditions`:** Conditions under which this step should be skipped (currently `None` for all steps).\n",
        "\n",
        "**Leveraging Gemini's Function Calling for Planning:**\n",
        "\n",
        "The `PlanningAgent` intelligently determines the need for steps like visualization and enhanced search (grounding) by utilizing Gemini's function calling capabilities. It analyzes the user's query and calls specific functions (e.g., `_determine_visualization_requirement`, `_determine_search_requirement`) to decide whether these steps are required. This dynamic plan creation based on query analysis demonstrates the power of combining structured planning with advanced language model features.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s0cyXQBfAel0"
      },
      "source": [
        "#### Query Validation and Suggestions\n",
        "\n",
        "A crucial part of the `PlanningAgent`'s role is to validate the user's query and provide suggestions for improvement. Let's see how this works:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X1AwbwP__pxh"
      },
      "outputs": [],
      "source": [
        "rich_print(plan.validated_query)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v9SEDT97HNYk"
      },
      "source": [
        "Here, the `PlanningAgent` has determined that the query is valid (`is_valid=True`) and has identified 'Austin' as the city of interest. It also confirms that no essential elements are missing (`missing_elements=[]`)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53RSiAZjAjiB"
      },
      "source": [
        "#### Query Suggestions\n",
        "\n",
        "Furthermore, the `PlanningAgent` provides suggestions to enhance the query:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rirBCWfS_5Dc"
      },
      "outputs": [],
      "source": [
        "rich_print(plan.validated_query.suggestions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B2FxhYZmDA8S"
      },
      "source": [
        "#### Failed Query\n",
        "\n",
        "What happens when the query is not valid? Let's see how the `PlanningAgent` handles such scenarios:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lad587v4J3P3"
      },
      "outputs": [],
      "source": [
        "agent = PlanningAgent(\n",
        "    query=\"I want to understand the EV charging situation in Paris and proper vetted information and some plot\",\n",
        "    client=client,\n",
        "    model_name=MODEL_ID_Flash,\n",
        "    debug=False,\n",
        ")\n",
        "plan = agent.create_plan()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uFEeB2a8IlrJ"
      },
      "source": [
        "In this case, the query mentions \"Paris,\" which is not a valid city in our predefined list (in `STATE_MAPPING`). The `PlanningAgent` detects this and returns:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U7YCYqosIz5C"
      },
      "outputs": [],
      "source": [
        "rich_print(plan.validated_query)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8BC-GeWPI3Ey"
      },
      "source": [
        "The `is_valid` flag is now `False`, and the `missing_elements` indicate that a \"valid city\" is required. Importantly, the `suggestions` provide specific guidance on how to correct the query, even suggesting valid city replacements."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4INx_rCJDjtp"
      },
      "outputs": [],
      "source": [
        "# You can see that it disable enabled search since the query didn't ask for anything \"enhance\" or \"grounding\"\n",
        "rich_print(plan.enable_search)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I9Z11rQwDfEE"
      },
      "outputs": [],
      "source": [
        "# it also skipped the visualization steps, since we didn't mention that in the query\n",
        "rich_print(plan.steps)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dYuYw015I8Du"
      },
      "source": [
        "Since the query was invalid, the `PlanningAgent` disables the search functionality (`enable_search=False`) and creates an empty list of steps (`steps=[]`). This effectively halts the execution process, as there's no valid plan to execute. This demonstrates the agent's ability to gracefully handle invalid queries and prevent unnecessary processing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FG322SGOGnrb"
      },
      "source": [
        "### Agent: QueryAnalysisAgent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lMeMdzKrKJAD"
      },
      "source": [
        "### Agent: QueryAnalysisAgent\n",
        "\n",
        "The `QueryAnalysisAgent` comes right after the `PlanningAgent` in our sequence. Its primary role is to dissect the user's query, identify key entities, and determine the type of analysis requested. It then passes this structured information to the next agent in the pipeline.\n",
        "\n",
        "**Input:**\n",
        "\n",
        "*   **Query:** The user's query about EV infrastructure, validated by the `PlanningAgent` (e.g., \"Analyze EV charging stations in Austin\").\n",
        "*   **Client:** An instance of the generative AI model client.\n",
        "*   **Model Name:** The specific model to be used (e.g., \"gemini-pro\").\n",
        "\n",
        "**Output:**\n",
        "\n",
        "*   **Dictionary:** Containing:\n",
        "    *   `status`: Whether the analysis was successful (\"success\" or \"error\").\n",
        "    *   `entities`: A dictionary representing the extracted entities from the query, based on the `QueryEntities` data model. This includes:\n",
        "        *   `pattern_type`: The type of analysis pattern detected (e.g., \"DISCOVERY\", \"COMPARISON\"). Although identified, these patterns are not yet used downstream in the current version but could be leveraged in future iterations.\n",
        "        *   `cities`: A list of valid cities extracted from the query.\n",
        "        *   `states`: A list of corresponding states for the extracted cities.\n",
        "        *   `research_theme`: The general theme of the query (currently fixed to \"Electronic Vehicle\").\n",
        "        *   `output_type`: The desired output type (e.g., \"Report\", \"Text\", \"Raw Data\").\n",
        "\n",
        "In essence, the `QueryAnalysisAgent` transforms the user's raw query into a structured format that can be easily understood and processed by the subsequent agents in the system. This section will delve into how the agent extracts these entities, handles different query patterns, and prepares the data for the next stage of the analysis."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SzUjpbDoGn8j"
      },
      "outputs": [],
      "source": [
        "from ev_agent.agent_handler.agent_03_QueryAnalysisAgent import *"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uq_s_HCWMVIL"
      },
      "source": [
        "Let's see how the `QueryAnalysisAgent` processes different types of queries.\n",
        "We'll examine three examples:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yQj--mhuqegX"
      },
      "source": [
        "#### Extraction Type 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dQM-sR_qMctq"
      },
      "source": [
        "Here, the query asks about gaps in Austin's charging network and requests a report format. The agent successfully analyzes the query and returns:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9bYBVxCoHF_S"
      },
      "outputs": [],
      "source": [
        "query = \"Where are the gaps in Austin charging network? Report format please\"\n",
        "query_agent = QueryAnalysisAgent(client, MODEL_ID)\n",
        "agent_1_result = query_agent.analyze(query)\n",
        "rich_print(agent_1_result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zGizyMvLMfEF"
      },
      "source": [
        "The agent correctly identifies the `pattern_type` as `GAPS`, extracts the city and state, and recognizes the desired `output_type` as `REPORT`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yOsBXcErqkLg"
      },
      "source": [
        "#### Extraction Type 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gxMepfZGMkYQ"
      },
      "source": [
        "In this case, the query requests raw data for Dallas. The agent responds with:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aaRoYZLUHbGL"
      },
      "outputs": [],
      "source": [
        "query = \"Need some raw data on Dallas for Ev charging stations\"\n",
        "agent_1_result = query_agent.analyze(query)\n",
        "rich_print(agent_1_result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kVWU7b4tMmnz"
      },
      "source": [
        "The agent identifies the `pattern_type` as `DISCOVERY` (since it's a general inquiry), extracts the city and state, and correctly sets the `output_type` to `RAW`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C5bWr2alqk8J"
      },
      "source": [
        "#### Extraction Type 3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ubffVGHoNMOY"
      },
      "outputs": [],
      "source": [
        "query = \"compare Dallas and Austin for EV Charging expansion and give me detail report.\"\n",
        "agent_1_result = query_agent.analyze(query)\n",
        "rich_print(agent_1_result)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cbM1lY83MvKT"
      },
      "source": [
        "These examples demonstrate the `QueryAnalysisAgent`'s ability to understand different query structures, extract relevant entities, and determine the user's intent regarding the analysis type and desired output format. This structured information is then passed on to subsequent agents in the pipeline, ensuring that the analysis stays focused and aligned with the user's needs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i-So9XlsE4L-"
      },
      "source": [
        "### Agent: DataGatherAgent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y4IaUnAqNSRu"
      },
      "source": [
        "### Agent: DataGatherAgent\n",
        "\n",
        "The `DataGatherAgent` is responsible for collecting the necessary data for our analysis by interacting with external APIs. It takes the structured output from the `QueryAnalysisAgent` and fetches relevant information about EV infrastructure and city demographics.\n",
        "\n",
        "**Input:**\n",
        "\n",
        "*   **`api_key`:**  Your NREL API key to access EV infrastructure data.\n",
        "*   **`radius_miles`:** The radius (in miles) around each city for which to gather data.\n",
        "*   **`debug`:** A boolean flag to enable/disable debug mode.\n",
        "\n",
        "**Output:**\n",
        "\n",
        "*   **`DataGatherAgentOutput`:** A data object containing:\n",
        "    *   `timestamp`: When the data was gathered.\n",
        "    *   `cities_data`: A list of `CityData` objects, one for each city in the query. Each `CityData` object may contain:\n",
        "        *   `city`: The name of the city.\n",
        "        *   `state`: The state abbreviation.\n",
        "        *   `summary`: A dictionary containing general city data retrieved from the Neighborhood Summary API (e.g., population, area, etc.).\n",
        "        *   `ev_data`: A dictionary containing EV charging station data retrieved from the EV Infrastructure Station Analysis API (e.g., number of stations, charger types, etc.).\n",
        "        *   `error`: Any error encountered while gathering data for the city.\n",
        "    *   `status`: The overall status of the data gathering process (\"success\" or \"error\").\n",
        "    *   `error`: Any general error encountered during the process.\n",
        "\n",
        "**Functionality:**\n",
        "\n",
        "The `DataGatherAgent` utilizes asynchronous programming (`asyncio`) to fetch data from two different APIs concurrently for each city:\n",
        "\n",
        "1. **Neighborhood Summary API:** Retrieves general demographic and infrastructure data about the city.\n",
        "2. **EV Infrastructure Station Analysis API:** Retrieves detailed information about EV charging stations within the specified radius.\n",
        "\n",
        "It handles potential errors during API calls, provides informative debug messages (if enabled), and compiles the gathered data into a structured `DataGatherAgentOutput` object. This agent plays a crucial role in bridging the gap between our analytical system and the real-world data needed to generate a meaningful report."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g9NKTXKHEpkq"
      },
      "outputs": [],
      "source": [
        "from ev_agent.agent_handler.agent_04_DataGatherAgent import *"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nAMEg6_COMq0"
      },
      "outputs": [],
      "source": [
        "# The Agent hits the OpenMapStreets API and NREL Developer API to gather data for a given city that can be helpful for Analysis.\n",
        "\n",
        "data_gather_agent = DataGatherAgent(\n",
        "    api_key=NREL_API_KEY, radius_miles=100.0, debug=False\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eTVSEplfOi5Y"
      },
      "source": [
        "Here, we create an instance of the `DataGatherAgent`, providing our `NREL_API_KEY`, a `radius_miles` of 100.00 miles, and setting `debug` to `True` to see detailed output."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tnsr3EzYOBqX"
      },
      "source": [
        "#### Single City"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tgpy8Q8-N9SM"
      },
      "outputs": [],
      "source": [
        "# Get the city from the QueryAnalysisAgent\n",
        "agent_1_result = query_agent.analyze(\n",
        "    \"Need some raw data on Dallas for Ev charging stations\"\n",
        ")\n",
        "\n",
        "# Get data from DataGatherAgent of the city\n",
        "agent_2_result = await data_gather_agent.process(agent_1_result)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XutOgx13ZRqr"
      },
      "outputs": [],
      "source": [
        "print(\"Number of cities given by the agent: \", len(agent_2_result.cities_data))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vj5qPtLJqotV"
      },
      "source": [
        "##### Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "15623699cf29"
      },
      "outputs": [],
      "source": [
        "# You can access the complete NeighborhoodSummary here:\n",
        "rich_print(\"NeighborhoodSummary - Complete \\n\", agent_2_result.cities_data[0].summary)\n",
        "\n",
        "\n",
        "# You can access the complete EVInfraSummary here:\n",
        "rich_print(\"EV Infra Summary - Complete \\n\", agent_2_result.cities_data[0].ev_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P4Kfv7uHOEsA"
      },
      "source": [
        "#### Multi City"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2OXo_jl-Epfz"
      },
      "outputs": [],
      "source": [
        "# Get the city from the QueryAnalysisAgent\n",
        "agent_1_result_multi_city = query_agent.analyze(\n",
        "    \"compare Dallas and Austin for EV Charging expansion and give me detail report\"\n",
        ")\n",
        "\n",
        "# Get data from DataGatherAgent of the city\n",
        "agent_2_result_multi_city = await data_gather_agent.process(agent_1_result_multi_city)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B66HlEAzZL1P"
      },
      "outputs": [],
      "source": [
        "print(\n",
        "    \"Number of cities given by the agent: \", len(agent_2_result_multi_city.cities_data)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bJgYsa0pqstF"
      },
      "source": [
        "##### Data - NeighborhoodSummary (OpenStreetMap - Overpass API)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wv0LiUHdsbi6"
      },
      "source": [
        "This API Handler uses Nomination API and Overpass API (OpenStreetMap). You can find more details [here](https://nominatim.org/), [here](https://nominatim.org/release-docs/develop/api/Overview/), [here](https://wiki.openstreetmap.org/wiki/Overpass_API)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "we3aopWOmiZj"
      },
      "outputs": [],
      "source": [
        "index = 0  # 0 for Dallas, 1 for Austin\n",
        "\n",
        "# You can see the NeighborhoodSummary of the city\n",
        "rich_print(\"City :\", agent_2_result_multi_city.cities_data[index].summary.city)\n",
        "rich_print(\"State :\", agent_2_result_multi_city.cities_data[index].summary.state)\n",
        "rich_print(\n",
        "    \"NeighborhoodSummary - Healthcare \\n\",\n",
        "    agent_2_result_multi_city.cities_data[index].summary.healthcare,\n",
        ")\n",
        "rich_print(\n",
        "    \"NeighborhoodSummary - Education \\n\",\n",
        "    agent_2_result_multi_city.cities_data[index].summary.education,\n",
        ")\n",
        "\n",
        "# You can see the complete data and all the elements of NeighborhoodSummary:\n",
        "rich_print(\n",
        "    \"NeighborhoodSummary - Complete \\n\",\n",
        "    agent_2_result_multi_city.cities_data[index].summary,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pVgA-XG2q-Q0"
      },
      "source": [
        "##### Data - EVInfraSummary (NREL Developer API)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OW5LGTu2ruGf"
      },
      "source": [
        "You can get more details about the API [here](https://developer.nrel.gov/) and [here](https://developer.nrel.gov/docs/transportation/alt-fuel-stations-v1/)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OfO4oJOvnK_o"
      },
      "outputs": [],
      "source": [
        "index = 0  # 0 for Dallas, 1 for Austin\n",
        "\n",
        "# You can see the EV Infra Summary of the city\n",
        "rich_print(\n",
        "    \"City :\", agent_2_result_multi_city.cities_data[index].ev_data.metadata[\"city\"]\n",
        ")\n",
        "rich_print(\n",
        "    \"State :\", agent_2_result_multi_city.cities_data[index].ev_data.metadata[\"state\"]\n",
        ")\n",
        "rich_print(\n",
        "    \"EV Infra Summary - Charging Capability \\n\",\n",
        "    agent_2_result_multi_city.cities_data[index].ev_data.charging_capabilities,\n",
        ")\n",
        "rich_print(\n",
        "    \"EV Infra Summary - Accessibility \\n\",\n",
        "    agent_2_result_multi_city.cities_data[index].ev_data.accessibility,\n",
        ")\n",
        "\n",
        "# You can see the complete data and all the elements of EV Infra Summary:\n",
        "# rich_print(\"EV Infra Summary - Complete \\n\",\n",
        "#            agent_2_result_multi_city.cities_data[index].ev_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x1fu8LdFqFR8"
      },
      "source": [
        "### Agent: ReportAgent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "93F5K5cjPjAn"
      },
      "source": [
        "### Agent: ReportAgent\n",
        "\n",
        "The `ReportAgent` takes the structured data gathered by the `DataGatherAgent` and transforms it into a comprehensive, well-formatted report. It's responsible for generating individual sections of the report, citing data sources appropriately, and optionally enhancing the content with information from web search.\n",
        "\n",
        "**Input:**\n",
        "\n",
        "*   **`client`:** An instance of the generative AI model client.\n",
        "*   **`model_name`:** The specific model to be used (e.g., \"gemini-pro-1.5\").\n",
        "*   **`enable_search`:** A boolean flag indicating whether to enhance the report with web search results.\n",
        "*   **`debug`:** A boolean flag to enable/disable debug mode.\n",
        "\n",
        "**Output:**\n",
        "\n",
        "*   **`Report`:** A data object containing the entire report, structured as follows:\n",
        "    *   `city`: The name of the city.\n",
        "    *   `state`: The state abbreviation.\n",
        "    *   `timestamp`: When the report was generated.\n",
        "    *   `sections`: A dictionary of `Section` objects, each representing a section of the report (e.g., \"Executive Summary\", \"Infrastructure Overview\"). Each `Section` includes:\n",
        "        *   `title`: The section title.\n",
        "        *   `content`: The main text content of the section.\n",
        "        *   `citations`: A dictionary of `CitationData` objects, mapping citation numbers to their corresponding data sources.\n",
        "        *   `key_findings`: A list of key takeaways from the section.\n",
        "        *   `enhanced_content`: Additional content generated through web search (if enabled).\n",
        "    *   `citations_text`: A formatted string containing all citations used in the report.\n",
        "    *   `full_text`: The entire report content as a single string.\n",
        "    *   `combined_report`: The full report content along with formatted citations.\n",
        "\n",
        "**Functionality:**\n",
        "\n",
        "The `ReportAgent` performs several key tasks:\n",
        "\n",
        "1. **Section Generation:** It generates individual report sections based on predefined templates and the gathered data, citing specific data points using a structured `CitationData` model.\n",
        "2. **Data Mapping:** It utilizes a detailed `_prepare_data_map` function to create a structured representation of the data from the `DataGatherAgent`, making it easier to reference specific data points in the report.\n",
        "3. **Asynchronous Processing:** It leverages asynchronous programming to generate multiple sections concurrently, improving efficiency.\n",
        "4. **Optional Search Enhancement:** If `enable_search` is set to `True`, it can enhance each section with information retrieved from Google Search, adding citations for the newly found data. This is achieved using the `_enhance_section_with_search` method.\n",
        "5. **Report Assembly:** Finally, it assembles the individual sections into a complete `Report` object, generating a formatted string representation of the entire report and its citations.\n",
        "\n",
        "The `ReportAgent` plays a critical role in synthesizing the raw data into a coherent, insightful, and well-supported analysis of the EV infrastructure. The following subsections will explore how this agent is used to generate reports, either for a single city with search grounding or for multiple cities without grounding."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ibu15GgxJ3G-"
      },
      "outputs": [],
      "source": [
        "from ev_agent.agent_handler.agent_05_ReportAgent import *"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pUiv3t3M7foo"
      },
      "source": [
        "#### Single City with grounding with Google"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ihv4IAfn-5gm"
      },
      "outputs": [],
      "source": [
        "report_agent_single_grounded = ReportAgent(\n",
        "    client=client, model_name=MODEL_ID_Flash, enable_search=True, debug=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TeqzU17t7ozN"
      },
      "outputs": [],
      "source": [
        "# Get the city from the QueryAnalysisAgent\n",
        "agent_1_result = query_agent.analyze(\n",
        "    \"Need some raw data on Dallas for Ev charging stations\"\n",
        ")\n",
        "rich_print(agent_1_result)\n",
        "\n",
        "# Get data from DataGatherAgent of the city\n",
        "agent_2_result = await data_gather_agent.process(agent_1_result)\n",
        "\n",
        "# Get the report built out using ReportAgent\n",
        "reports_single_grounded = await report_agent_single_grounded.analyze(\n",
        "    agent_1_result, agent_2_result\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UCmgBSuU7SQC"
      },
      "outputs": [],
      "source": [
        "print(\n",
        "    \"Report is on the city: \",\n",
        "    reports_single_grounded.city,\n",
        "    \" and state: \",\n",
        "    reports_single_grounded.state,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "biVMqgbuAwMY"
      },
      "source": [
        "Predefined/Available Section of the Reports:\n",
        "\n",
        "*   Executive Summary\n",
        "*   Infrastructure Overview\n",
        "*   Current EV Assessment\n",
        "*   Demand Analysis\n",
        "*   Supply Analysis\n",
        "*   Gap Analysis\n",
        "*   Location Recommendations\n",
        "*   Implementation Strategy\n",
        "\n",
        "You can explore each section and see how grounding with Google, enhanced the section with updated text and citations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-Q7T-w9R7hwd"
      },
      "outputs": [],
      "source": [
        "for section_name, section_text in reports_single_grounded.sections.items():\n",
        "    if section_name == \"Infrastructure Overview\":\n",
        "        print(section_name)\n",
        "        rich_print(section_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c3d5741c77ca"
      },
      "outputs": [],
      "source": [
        "# You can access other key areas of the report:\n",
        "reports_single_grounded.full_text  # Full text of the report - without citations\n",
        "reports_single_grounded.citations_text  # Full text of the citations - without text\n",
        "reports_single_grounded.combined_report  # Full text of the report combined with citations\n",
        "reports_single_grounded.timestamp  # Timestamp of report generations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vt9f7S5v7icz"
      },
      "source": [
        "#### Multi City without grounding with Google"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Ifsrp-G-_T1"
      },
      "outputs": [],
      "source": [
        "report_agent_multi_city = ReportAgent(\n",
        "    client=client,\n",
        "    model_name=MODEL_ID_Flash,\n",
        "    enable_search=False,  # you can enable grounding for both the cities if you want\n",
        "    debug=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TpRaJx-v7kNm"
      },
      "outputs": [],
      "source": [
        "# Get the city from the QueryAnalysisAgent\n",
        "agent_1_result_multi_city = query_agent.analyze(\n",
        "    \"compare Dallas and Austin for EV Charging expansion and give me detail report\"\n",
        ")\n",
        "rich_print(agent_1_result_multi_city)\n",
        "\n",
        "# Get data from DataGatherAgent of the city\n",
        "agent_2_result_multi_city = await data_gather_agent.process(agent_1_result_multi_city)\n",
        "\n",
        "# Get the report built out using ReportAgent\n",
        "reports_multi_city = await report_agent_multi_city.analyze(\n",
        "    agent_1_result_multi_city, agent_2_result_multi_city\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ph40v5rp7wPK"
      },
      "outputs": [],
      "source": [
        "index = 0\n",
        "print(\n",
        "    \"Report is on the city: \",\n",
        "    reports_multi_city[index].city,\n",
        "    \" and state: \",\n",
        "    reports_multi_city[index].state,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zn1LtOvi7wIU"
      },
      "outputs": [],
      "source": [
        "index = 1\n",
        "print(\n",
        "    \"Report is on the city: \",\n",
        "    reports_multi_city[index].city,\n",
        "    \" and state: \",\n",
        "    reports_multi_city[index].state,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kmU1fJFB8PTi"
      },
      "outputs": [],
      "source": [
        "for section_name, section_text in reports_multi_city[index].sections.items():\n",
        "    if section_name == \"Demand Analysis\":\n",
        "        print(section_name)\n",
        "        rich_print(section_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "582a01dfc15d"
      },
      "outputs": [],
      "source": [
        "# You can also check all the sections using object\n",
        "\n",
        "print(\"All the sections in the report\")\n",
        "for section_name, section_text in reports_multi_city[index].sections.items():\n",
        "    print(section_name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3e0df2853b20"
      },
      "outputs": [],
      "source": [
        "# You can access other key areas of the report by passing appropriate indexes:\n",
        "\n",
        "reports_multi_city[index].full_text  # Full text of the report - without citations\n",
        "\n",
        "reports_multi_city[index].citations_text  # Full text of the citations - without text\n",
        "\n",
        "reports_multi_city[\n",
        "    index\n",
        "].combined_report  # Full text of the report combined with citations\n",
        "\n",
        "reports_multi_city[index].timestamp  # Timestamp of report generations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Yd9DWWLP7T2"
      },
      "source": [
        "### Agent: VisualizeAgent\n",
        "\n",
        "The `VisualizeAgent` is responsible for creating insightful visualizations based on the data gathered by the `DataGatherAgent`. It uses the `plotly` library to generate various charts and graphs that help to understand the EV infrastructure landscape in a more visual and intuitive manner. Although it's called an \"agent\" here, it's important to note that this is essentially a set of helper functions for creating visualizations rather than an autonomous agent with decision-making capabilities.\n",
        "\n",
        "**Input:**\n",
        "\n",
        "*   **`data`:** The `DataGatherAgentOutput` object, containing structured data for one or more cities.\n",
        "\n",
        "**Output:**\n",
        "\n",
        "*   A tuple containing two dictionaries:\n",
        "    *   **`single_city_figs`:** A dictionary of `plotly` figure objects, each representing a visualization specific to a single city.\n",
        "    *   **`comparison_figs`:** A dictionary of `plotly` figure objects, each representing a comparative visualization across multiple cities (if applicable).\n",
        "\n",
        "**Functionality:**\n",
        "\n",
        "The `VisualizeAgent` performs the following tasks:\n",
        "\n",
        "1. **Single City Visualizations:** It generates a set of visualizations for each city using the `create_comprehensive_city_analysis` function. These include:\n",
        "    *   **EV Infrastructure Overview:** Bar charts showing charging station types, connector distribution, network distribution, and access & payment methods.\n",
        "    *   **Transportation Infrastructure Analysis:** A multi-panel plot showing public transport facilities, road network distribution, parking facilities, and a comparison of EV vs. traditional vehicle infrastructure.\n",
        "    *   **Urban Amenities and Services:** A multi-panel plot showing the distribution of retail and shopping centers, food and entertainment venues, emergency services, and public amenities.\n",
        "    *   **Area Analysis:** A pie chart displaying the distribution of total area, water area, green area, and built area.\n",
        "\n",
        "2. **Multi-City Comparisons (if applicable):** If the input data contains information for multiple cities, it uses the `plot_multi_city_comparison` function to generate comparative visualizations. These include:\n",
        "    *   **EV Infrastructure Comparisons:** Bar charts comparing the number of EV stations vs. fuel stations, charging station types, and EV station density across cities.\n",
        "    *   **Transportation Infrastructure:** Bar charts comparing public transport infrastructure, road network distribution, and parking facilities across cities.\n",
        "    *   **Area Analysis:** A bar chart comparing area distribution (total, water, green, built) across cities.\n",
        "    *   **Urban Amenities:** A bar chart comparing the prevalence of various urban amenities (e.g., shopping centers, restaurants, hospitals) across cities.\n",
        "\n",
        "3. **Visualization Organization:** It organizes all generated plots into the `single_city_figs` and `comparison_figs` dictionaries, making it easy to access specific visualizations.\n",
        "\n",
        "The `VisualizeAgent` plays a crucial role in making the data more accessible and understandable by providing a visual representation of key metrics and trends. These visualizations can aid in identifying patterns, making comparisons, and ultimately supporting decision-making related to EV infrastructure planning and development."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LqGemlxj-ohw"
      },
      "outputs": [],
      "source": [
        "from ev_agent.agent_handler.agent_06_VisualizeAgent import *"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ly_pvvoNCKr_"
      },
      "source": [
        "#### Single City"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zjveHtOhGzfn"
      },
      "outputs": [],
      "source": [
        "single_city_figs, comparison_figs = plot_all_visualizations(agent_2_result)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SM9nbe5FHEap"
      },
      "outputs": [],
      "source": [
        "print(\"\\n=== Single City Analysis ===\")\n",
        "for name, fig in single_city_figs.items():\n",
        "    print(f\"\\nDisplaying: {name.replace('_', ' ').title()}\")\n",
        "    fig.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5iF_b6oMCVBL"
      },
      "source": [
        "#### Multi City"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "33fASgQGHP6d"
      },
      "outputs": [],
      "source": [
        "single_city_figs, comparison_figs = plot_all_visualizations(agent_2_result_multi_city)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dZuNdJeoHTVD"
      },
      "outputs": [],
      "source": [
        "print(\"\\n=== Multi-City Comparisons ===\")\n",
        "for name, fig in comparison_figs.items():\n",
        "    print(f\"\\nDisplaying: {name.replace('_', ' ').title()}\")\n",
        "    fig.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o4cW5bVtQ2Lo"
      },
      "source": [
        "## Next Steps and Potential Improvements\n",
        "\n",
        "We've built a solid foundation for a multi-agent system that analyzes EV infrastructure. However, there's always room for improvement and expansion. Here are some potential next steps, inspired by features found in advanced multi-agent frameworks like AutoGen, CrewAI, and LangGraph:\n",
        "\n",
        "1. **Enhanced Agent Communication:** Implement dynamic inter-agent communication for iterative feedback, dynamic task allocation, and agent specialization.\n",
        "2. **Sophisticated Planning:** Develop more advanced planning with conditional logic, sub-planning, and plan repair capabilities.\n",
        "3. **Expanded Tool Integration:** Integrate with more APIs, databases, web scraping, and knowledge graphs to broaden the system's knowledge base.\n",
        "4. **Interactive User Experience:** Allow for clarification dialogs, progress updates, interactive visualizations, and user feedback mechanisms.\n",
        "5. **Robust Error Handling:** Implement comprehensive exception handling, retry mechanisms, and fallback strategies for increased reliability.\n",
        "6. **Integrated Visualizations:** Incorporate visualizations directly into the generated reports for a more cohesive and engaging presentation.\n",
        "7. **Agent Memory and Learning:** Introduce agent memory for caching, learning from user feedback, and potential model fine-tuning to improve performance over time.\n",
        "\n",
        "By implementing these enhancements, we can create a more powerful, flexible, and user-friendly multi-agent system for analyzing EV infrastructure and generating actionable insights."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "intro_research_multi_agents_gemini_2_0.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}