{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ijGzTHJJUCPY"
      },
      "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": "r11Gu7qNgx1p"
      },
      "source": [
        "# Building a photo recognition agent: Agent Engine setup\n",
        "\n",
        "This notebook shows how to setup the Agent Engine for the photo recognition agent demo"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK for Python\n",
        "\n",
        "Install the latest version of the Vertex AI SDK for Python as well as extra dependencies related to Agent Engine and LangChain:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --quiet google-cloud-aiplatform[langchain,agent_engines]==1.83.0 google-cloud-discoveryengine==0.11.10"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart current 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 will restart the current kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "# Restart kernel after installs so that your environment can access the new packages\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\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": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"<YOUR GOOGLE CLOUD PROJECT ID>\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "STAGING_BUCKET = \"gs://<YOUR GCS BUCKET>\"  # @param {type:\"string\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gnGK_0Du4LD1"
      },
      "source": [
        "## Wikipedia Tool\n",
        "In this section, you'll define multiple Python functions that access the Wikipedia API. Later we'll use these Python functions as Tools in our Agent Engine agent."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TOJCaEOF4LD1"
      },
      "source": [
        "### Define functions for accessing the Wikipedia API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "go8KGzoa4LD1"
      },
      "outputs": [],
      "source": [
        "import requests\n",
        "\n",
        "\n",
        "# search for a wiki page\n",
        "def search_wiki_title(query):\n",
        "    url = f\"https://en.wikipedia.org/w/api.php?action=query&list=search&srsearch={query}&srlimit=1&format=json\"\n",
        "    resp = requests.get(url)\n",
        "    return resp.json()[\"query\"][\"search\"][0][\"title\"]\n",
        "\n",
        "\n",
        "# get full text of the wiki page\n",
        "def get_wiki_full_text(wiki_title):\n",
        "    # get the page\n",
        "    url = f\"https://en.wikipedia.org/w/api.php?action=query&prop=extracts&titles={wiki_title}&explaintext=true&format=json\"\n",
        "    response = requests.get(url)\n",
        "    data = response.json()\n",
        "\n",
        "    # extract plain text\n",
        "    page_id = next(iter(data[\"query\"][\"pages\"]))\n",
        "    plain_text = data[\"query\"][\"pages\"][page_id][\"extract\"]\n",
        "    return plain_text"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-YlC37yz4LD1"
      },
      "source": [
        "### Define Tool function for the Wikipedia API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QfLgfINE4LD1"
      },
      "outputs": [],
      "source": [
        "# make a query on Wikipedia for a topic\n",
        "\n",
        "\n",
        "def query_with_wikipedia(\n",
        "    query: str = \"Fallingwater\",\n",
        "):\n",
        "    \"\"\"\n",
        "    Finds answer for any topic or object from Wikipedia and returns a dictionary containing the answer.\n",
        "\n",
        "    Args:\n",
        "        query: the name of object or topic to find.\n",
        "\n",
        "    Example: {\"answer\": \"Fallingwater is a house designed by the architect Frank Lloyd Wright in 1935.\"}\n",
        "    \"\"\"\n",
        "\n",
        "    wiki_title = search_wiki_title(query)  # calls Wikipedia API to get wiki page title\n",
        "    wiki_full_text = get_wiki_full_text(\n",
        "        wiki_title\n",
        "    )  # calls Wikipedia API to get full text\n",
        "\n",
        "    return {\"answer\": wiki_full_text}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m5Ec-fkg4LD1"
      },
      "outputs": [],
      "source": [
        "query_with_wikipedia()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qkjkFJHg4LD1"
      },
      "source": [
        "## Vertex AI Search Tool\n",
        "In this section, it will define a Tool for calling the /ask_gms on the Cloud Run for making a query with the Google Merch Shop dataset on Vertex AI Search. So, before using this Tool, you need to deploy the Run instance with /ag-web/app/deploy.sh ."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oCu-ih3W4LD2"
      },
      "source": [
        "### Define Tool function for Vertex AI Search\n",
        "\n",
        "The Cloud Run host name can be found on Console > Cloud Run > ag-web > URL at the top"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A2x4OEIg4LD2"
      },
      "outputs": [],
      "source": [
        "from vertexai import agent_engines\n",
        "from vertexai.preview import reasoning_engines\n",
        "\n",
        "GOOGLE_SHOP_VERTEXAI_SEARCH_URL = (\n",
        "    \"https://ag-web-nhhfh7g7iq-uc.a.run.app/ask_gms\"  # please change\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7nPD25EK4LD2"
      },
      "outputs": [],
      "source": [
        "def find_product_from_googleshop(product_name: str, product_description: str):\n",
        "    \"\"\"\n",
        "    Find a product with the product_name and product_description from\n",
        "    Google Merch Shop and returns a dictionary containing product details.\n",
        "    \"\"\"\n",
        "\n",
        "    params = {\"query\": product_name + \" \" + product_description}\n",
        "    response = requests.get(\n",
        "        GOOGLE_SHOP_VERTEXAI_SEARCH_URL, params\n",
        "    )  # calls Vertex AI Search\n",
        "    item = response.json()\n",
        "    productDetails = f\"\"\"\n",
        "        {item['gms_name']} is a product sold at Google Merch Shop. The price is {item['price']}.\n",
        "        {item['gms_desc']}. You can buy the product at their web site: {item['link']}\"\n",
        "    \"\"\"\n",
        "\n",
        "    return {\"productDetails\": productDetails}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xVjNIewU4LD2"
      },
      "source": [
        "## Define and deploy the Agent\n",
        "In this section, it creates an agent with the Wikipedia Tool and Search Tool, and deploy it to the Agent Engine runtime.\n",
        "\n",
        "### Define the agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zbJXWoiZ4LD2"
      },
      "outputs": [],
      "source": [
        "model_name = \"gemini-2.0-flash\"\n",
        "agent = reasoning_engines.LangchainAgent(\n",
        "    model=model_name,\n",
        "    tools=[query_with_wikipedia, find_product_from_googleshop],\n",
        "    agent_executor_kwargs={\"return_intermediate_steps\": True},\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bv8YYxGH4LD2"
      },
      "source": [
        "### Test it locally"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TylY4MDX4LD2"
      },
      "outputs": [],
      "source": [
        "# input_text = \"Google Dino Pin. A collectible enamel pin featuring the iconic Lonely T-Rex dinosaur from the Google Chrome browser offline error page. Great product.\"\n",
        "# input_text = \"Google Bike Enamel Pin. This is a collectible lapel pin featuring a colorful bicycle with a basket, likely representing Google's company culture and its promotion of cycling as a sustainable mode of transportation.\"\n",
        "# input_text = \"Android Pen. This is a blue ballpoint pen featuring the Android logo. It appears to be a promotional or branded item associated with the Android operating system\"\n",
        "# input_text = \"Google Chrome Dinosaur Game Pin. where can i buy it? Write the answer in plaintext.\"\n",
        "input_text = (\n",
        "    \"Google Bike Enamel Pin. where can i buy it? Write the answer in plaintext.\"\n",
        ")\n",
        "agent.query(input=input_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gpz4arAz4LD2"
      },
      "source": [
        "### Deploy the Agent to the Agent Engine runtime"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r1kDkwQk4LD2"
      },
      "outputs": [],
      "source": [
        "remote_agent = agent_engines.create(\n",
        "    agent,\n",
        "    requirements=[\n",
        "        \"google-cloud-aiplatform[langchain,agent_engines]\",\n",
        "        \"requests\",\n",
        "    ],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q_8ZKk224LD2"
      },
      "source": [
        "### Test it"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hu5J03O_4LD2"
      },
      "outputs": [],
      "source": [
        "remote_agent.query(input=\"What is Google Bike Enamel Pin?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ubi3lMBU4LD2"
      },
      "outputs": [],
      "source": [
        "remote_agent.query(input=\"What is Fallingwater?\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "ag_setup_re.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
