{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qDHHyJsZdFFp"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 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": "THTB3L8TxZ1Q"
      },
      "source": [
        "# Getting Started With MCP Toolbox\n",
        "\n",
        "This guide demonstrates how to quickly run\n",
        "[Toolbox](https://github.com/googleapis/genai-toolbox) end-to-end in Google\n",
        "Colab using Python, BigQuery, and either [Google\n",
        "GenAI](https://pypi.org/project/google-genai/), [ADK](https://google.github.io/adk-docs/),\n",
        "[Langgraph](https://www.langchain.com/langgraph)\n",
        "or [LlamaIndex](https://www.llamaindex.ai/).\n",
        "\n",
        "Within this Colab environment, you'll\n",
        "- Set up a `BigQuery Dataset`.\n",
        "- Launch a Toolbox server.\n",
        "- Connect to Toolbox and develop a sample `Hotel Booking` application."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KLQzss0WxeI1"
      },
      "source": [
        "## Step 1: Set up your dataset\n",
        "\n",
        "In this section, we will\n",
        "1. Create a dataset in your bigquery project.\n",
        "1. Insert example data into the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zTtKdvbwAag3"
      },
      "outputs": [],
      "source": [
        "# @markdown Please fill in the value below and then run the cell.\n",
        "BIGQUERY_PROJECT = \"\" # @param {type:\"string\"}\n",
        "DATASET = \"toolbox_ds\" # @param {type:\"string\"}\n",
        "TABLE_ID = \"hotels\" # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bDaRyfx3PhXM"
      },
      "source": [
        "> You need to authenticate as an IAM user so this notebook can access your Google Cloud Project. This access is necessary to use Google's LLM models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c1_GR5NwPhXM"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth\n",
        "\n",
        "# Authenticate the user for Google Cloud access\n",
        "auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2eNdr9LYyhuV",
        "outputId": "4ded8803-5b9c-4a26-af03-28a6f23a415e"
      },
      "outputs": [],
      "source": [
        "# Create the dataset if it does not exist\n",
        "from google.cloud import bigquery\n",
        "from google.cloud import exceptions\n",
        "\n",
        "bqclient = bigquery.Client(project=BIGQUERY_PROJECT)\n",
        "dataset_ref = bqclient.dataset(DATASET)\n",
        "\n",
        "# Check if the dataset already exists\n",
        "try:\n",
        "    bqclient.get_dataset(dataset_ref)\n",
        "    print(f\"Dataset {DATASET} already exists. Skipping creation.\")\n",
        "except exceptions.NotFound:\n",
        "    # If a google.cloud.exceptions.NotFound error is raised, the dataset does not exist.\n",
        "    print(f\"Dataset {DATASET} not found. Creating dataset...\")\n",
        "    dataset = bigquery.Dataset(dataset_ref)\n",
        "    dataset = bqclient.create_dataset(dataset)\n",
        "    print(f\"Dataset {DATASET} created successfully.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 432
        },
        "id": "6t_nLJIHCRgy",
        "outputId": "11a5c025-288d-43e2-a823-454b0d88b1ce"
      },
      "outputs": [],
      "source": [
        "table_ref = dataset_ref.table(TABLE_ID)\n",
        "\n",
        "schema = [\n",
        "    bigquery.SchemaField(\"id\", \"INTEGER\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"name\", \"STRING\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"location\", \"STRING\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"price_tier\", \"STRING\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"checkin_date\", \"DATE\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"checkout_date\", \"DATE\", mode=\"REQUIRED\"),\n",
        "    bigquery.SchemaField(\"booked\", \"BOOLEAN\", mode=\"REQUIRED\"),\n",
        "]\n",
        "\n",
        "# Check if the table already exists; if not, create it and insert data\n",
        "try:\n",
        "    bqclient.get_table(table_ref)\n",
        "    raise ValueError(f\"Table '{TABLE_ID}' already exists in dataset '{DATASET}'. Please delete it or use a different table name.\")\n",
        "except exceptions.NotFound:\n",
        "    table = bigquery.Table(table_ref, schema=schema)\n",
        "    table = bqclient.create_table(table)\n",
        "    print(f\"Created table '{TABLE_ID}'.\")\n",
        "\n",
        "    sql = f\"\"\"\n",
        "    INSERT INTO `{BIGQUERY_PROJECT}.{DATASET}.{TABLE_ID}`(id, name, location, price_tier, checkin_date, checkout_date, booked)\n",
        "    VALUES\n",
        "      (1, 'Hilton Basel', 'Basel', 'Luxury', '2024-04-20', '2024-04-22', FALSE),\n",
        "      (2, 'Marriott Zurich', 'Zurich', 'Upscale', '2024-04-14', '2024-04-21', FALSE),\n",
        "      (3, 'Hyatt Regency Basel', 'Basel', 'Upper Upscale', '2024-04-02', '2024-04-20', FALSE),\n",
        "      (4, 'Radisson Blu Lucerne', 'Lucerne', 'Midscale', '2024-04-05', '2024-04-24', FALSE),\n",
        "      (5, 'Best Western Bern', 'Bern', 'Upper Midscale', '2024-04-01', '2024-04-23', FALSE),\n",
        "      (6, 'InterContinental Geneva', 'Geneva', 'Luxury', '2024-04-23', '2024-04-28', FALSE),\n",
        "      (7, 'Sheraton Zurich', 'Zurich', 'Upper Upscale', '2024-04-02', '2024-04-27', FALSE),\n",
        "      (8, 'Holiday Inn Basel', 'Basel', 'Upper Midscale', '2024-04-09', '2024-04-24', FALSE),\n",
        "      (9, 'Courtyard Zurich', 'Zurich', 'Upscale', '2024-04-03', '2024-04-13', FALSE),\n",
        "      (10, 'Comfort Inn Bern', 'Bern', 'Midscale', '2024-04-04', '2024-04-16', FALSE);\n",
        "    \"\"\"\n",
        "    job = bqclient.query(sql)\n",
        "    job.result()\n",
        "    print(\"Data inserted into 'hotels' table.\")\n",
        "\n",
        "sql_select = f\"SELECT * FROM `{BIGQUERY_PROJECT}.{DATASET}.{TABLE_ID}`\"\n",
        "query_job = bqclient.query(sql_select)\n",
        "\n",
        "print(\"\\nTable Content:\")\n",
        "query_job.to_dataframe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EPuheP8DIt3p"
      },
      "source": [
        "## Step 2: Install and configure Toolbox\n",
        "\n",
        "In this section, we will\n",
        "1. Download the latest version of the toolbox binary.\n",
        "2. Create a toolbox config file.\n",
        "3. Start a toolbox server using the config file.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bl1IeaqZbMYh"
      },
      "source": [
        "Download the [latest](https://github.com/googleapis/genai-toolbox/releases) version of Toolbox as a binary."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lbsQ1Aa-IszB",
        "outputId": "07c57730-b285-4069-e6e1-97d128cdcda4"
      },
      "outputs": [],
      "source": [
        "version = \"0.25.0\" # x-release-please-version\n",
        "! curl -O https://storage.googleapis.com/genai-toolbox/v{version}/linux/amd64/toolbox\n",
        "\n",
        "# Make the binary executable\n",
        "! chmod +x toolbox"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ovlzi2RVJGM5"
      },
      "outputs": [],
      "source": [
        "TOOLBOX_BINARY_PATH = \"/content/toolbox\"\n",
        "SERVER_PORT = 5000"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KNg7v_FeTYJu"
      },
      "source": [
        "Create a tools file with the following functions:\n",
        "\n",
        "- `Database Connection (sources)`: `Includes details for connecting to our hotels database.`\n",
        "- `Tool Definitions (tools)`: `Defines five tools for database interaction:`\n",
        "  - `search-hotels-by-name`\n",
        "  - `search-hotels-by-location`\n",
        "  - `book-hotel`\n",
        "  - `update-hotel`\n",
        "  - `cancel-hotel`\n",
        "\n",
        "Our application will leverage these tools to interact with the hotels table.\n",
        "\n",
        "For detailed configuration options, please refer to the [Toolbox documentation](https://googleapis.github.io/genai-toolbox/getting-started/configure/).\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jje8N5fScchw"
      },
      "outputs": [],
      "source": [
        "# Create a tools file at runtime.\n",
        "# You can also upload a tools file and use that to run toolbox.\n",
        "tools_file_name = \"tools.yml\"\n",
        "file_content = f\"\"\"\n",
        "sources:\n",
        "  my-bigquery-source:\n",
        "    kind: bigquery\n",
        "    project: {BIGQUERY_PROJECT}\n",
        "tools:\n",
        "  search-hotels-by-name:\n",
        "    kind: bigquery-sql\n",
        "    source: my-bigquery-source\n",
        "    description: Search for hotels based on name.\n",
        "    parameters:\n",
        "      - name: name\n",
        "        type: string\n",
        "        description: The name of the hotel.\n",
        "    statement: SELECT * FROM `{DATASET}.{TABLE_ID}` WHERE LOWER(name) LIKE LOWER(CONCAT('%', @name, '%'));\n",
        "  search-hotels-by-location:\n",
        "    kind: bigquery-sql\n",
        "    source: my-bigquery-source\n",
        "    description: Search for hotels based on location.\n",
        "    parameters:\n",
        "      - name: location\n",
        "        type: string\n",
        "        description: The location of the hotel.\n",
        "    statement: SELECT * FROM `{DATASET}.{TABLE_ID}` WHERE LOWER(location) LIKE LOWER(CONCAT('%', @location, '%'));\n",
        "  book-hotel:\n",
        "    kind: bigquery-sql\n",
        "    source: my-bigquery-source\n",
        "    description: >-\n",
        "       Book a hotel by its ID. If the hotel is successfully booked, returns a NULL, raises an error if not.\n",
        "    parameters:\n",
        "      - name: hotel_id\n",
        "        type: integer\n",
        "        description: The ID of the hotel to book.\n",
        "    statement: UPDATE `{DATASET}.{TABLE_ID}` SET booked = TRUE WHERE id = @hotel_id;\n",
        "  update-hotel:\n",
        "    kind: bigquery-sql\n",
        "    source: my-bigquery-source\n",
        "    description: >-\n",
        "      Update a hotel's check-in and check-out dates by its ID. Returns a message indicating whether the hotel was successfully updated or not.\n",
        "    parameters:\n",
        "      - name: checkin_date\n",
        "        type: string\n",
        "        description: The new check-in date of the hotel.\n",
        "      - name: checkout_date\n",
        "        type: string\n",
        "        description: The new check-out date of the hotel.\n",
        "      - name: hotel_id\n",
        "        type: integer\n",
        "        description: The ID of the hotel to update.\n",
        "    statement: >-\n",
        "      UPDATE `{DATASET}.{TABLE_ID}` SET checkin_date = PARSE_DATE('%Y-%m-%d', @checkin_date), checkout_date = PARSE_DATE('%Y-%m-%d', @checkout_date) WHERE id = @hotel_id;\n",
        "  cancel-hotel:\n",
        "    kind: bigquery-sql\n",
        "    source: my-bigquery-source\n",
        "    description: Cancel a hotel by its ID.\n",
        "    parameters:\n",
        "      - name: hotel_id\n",
        "        type: integer\n",
        "        description: The ID of the hotel to cancel.\n",
        "    statement: UPDATE `{DATASET}.{TABLE_ID}` SET booked = FALSE WHERE id = @hotel_id;\n",
        "toolsets:\n",
        "  my-toolset:\n",
        "    - search-hotels-by-name\n",
        "    - search-hotels-by-location\n",
        "    - book-hotel\n",
        "    - update-hotel\n",
        "    - cancel-hotel\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JPNXr4y58tMH"
      },
      "outputs": [],
      "source": [
        "with open(tools_file_name, 'w', encoding='utf-8') as f:\n",
        "    f.write(file_content)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5ZH5VuYzdP_W"
      },
      "outputs": [],
      "source": [
        "TOOLS_FILE_PATH = f\"/content/{tools_file_name}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iZGQzYUF-pho"
      },
      "outputs": [],
      "source": [
        "# Start a toolbox server\n",
        "! nohup {TOOLBOX_BINARY_PATH} --tools-file {TOOLS_FILE_PATH} -p {SERVER_PORT} > toolbox.log 2>&1 &"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1PJpKOBieKOV",
        "outputId": "3bc866bd-86b1-4b5c-f77a-f0acf65ebd4e"
      },
      "outputs": [],
      "source": [
        "# Check if toolbox is running\n",
        "!sudo lsof -i :{SERVER_PORT}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4yFH4JK7JEAv"
      },
      "source": [
        "## Step 3: Connect your agent to Toolbox\n",
        "\n",
        "In this section, you will\n",
        "1. Establish a connection to the tools by creating a Toolbox client.\n",
        "2. Build an agent that leverages the tools and an LLM for Hotel Booking functionality.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u0Jc-0YNdhQd",
        "outputId": "8c0167ca-bcca-492a-e7b6-aab808ecc156"
      },
      "outputs": [],
      "source": [
        "# Configure gcloud.\n",
        "!gcloud config set project {BIGQUERY_PROJECT}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J46eLkFbNhWq"
      },
      "source": [
        "> You can use ADK, LangGraph, or LlamaIndex to develop a Toolbox based application. Run one of the [Connect Using LangGraph](#scrollTo=pbapNMhhL33S), [Connect using LlamaIndex](#scrollTo=04iysrm_L_7v&line=1&uniqifier=1) or [Connect using ADK](#scrollTo=yA3rAiELIds5) sections below.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pbapNMhhL33S"
      },
      "source": [
        "### Connect Using LangGraph"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uraBx8mbMXnV",
        "outputId": "d4fc8823-e8e5-4f55-95e8-06eafbb5e0b1"
      },
      "outputs": [],
      "source": [
        "# Install the Toolbox Langchain package\n",
        "!pip install toolbox-langchain --quiet\n",
        "!pip install langgraph --quiet\n",
        "\n",
        "# Install the Langchain llm package\n",
        "# TODO(developer): replace this with another model if needed\n",
        "! pip install langchain-google-vertexai --quiet\n",
        "# ! pip install langchain-google-genai\n",
        "# ! pip install langchain-anthropic"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0oHNnZnBM8FU"
      },
      "source": [
        "Create a LangGraph Hotel Agent which can Search, Book and Cancel hotels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Br3ucM46M9uc",
        "outputId": "7118993f-d5f7-4e15-ba28-0dc71cc6c403"
      },
      "outputs": [],
      "source": [
        "from langgraph.prebuilt import create_react_agent\n",
        "# TODO(developer): replace this with another import if needed\n",
        "from langchain_google_vertexai import ChatVertexAI\n",
        "# from langchain_google_genai import ChatGoogleGenerativeAI\n",
        "# from langchain_anthropic import ChatAnthropic\n",
        "from langgraph.checkpoint.memory import MemorySaver\n",
        "\n",
        "from toolbox_langchain import ToolboxClient\n",
        "\n",
        "prompt = \"\"\"\n",
        "  You're a helpful hotel assistant. You handle hotel searching, booking and\n",
        "  cancellations. When the user searches for a hotel, mention it's name, id,\n",
        "  location and price tier. Always mention hotel id while performing any\n",
        "  searches. This is very important for any operations. For any bookings or\n",
        "  cancellations, please provide the appropriate confirmation. Be sure to\n",
        "  update checkin or checkout dates if mentioned by the user.\n",
        "  Don't ask for confirmations from the user.\n",
        "\"\"\"\n",
        "\n",
        "queries = [\n",
        "    \"Find hotels in Basel with Basel in it's name.\",\n",
        "    \"Can you book the Hilton Basel for me?\",\n",
        "    \"Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.\",\n",
        "    \"My check in dates would be from April 10, 2024 to April 19, 2024.\",\n",
        "]\n",
        "\n",
        "# Create an LLM to bind with the agent.\n",
        "# TODO(developer): replace this with another model if needed\n",
        "model = ChatVertexAI(model_name=\"gemini-2.0-flash-001\", project=BIGQUERY_PROJECT)\n",
        "# model = ChatGoogleGenerativeAI(model=\"gemini-2.0-flash-001\")\n",
        "# model = ChatAnthropic(model=\"claude-3-5-sonnet-20240620\")\n",
        "\n",
        "# Load the tools from the Toolbox server\n",
        "client = ToolboxClient(\"http://127.0.0.1:5000\")\n",
        "tools = client.load_toolset()\n",
        "\n",
        "# Create a Langraph agent\n",
        "agent = create_react_agent(model, tools, checkpointer=MemorySaver())\n",
        "config = {\"configurable\": {\"thread_id\": \"thread-1\"}}\n",
        "for query in queries:\n",
        "    inputs = {\"messages\": [(\"user\", prompt + query)]}\n",
        "    response = agent.invoke(inputs, stream_mode=\"values\", config=config)\n",
        "    print(response[\"messages\"][-1].content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04iysrm_L_7v"
      },
      "source": [
        "### Connect using LlamaIndex"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6b6Loh8SJ_iA",
        "outputId": "a834e07b-6f28-48f2-ef26-1204af5e053d"
      },
      "outputs": [],
      "source": [
        "# Install the Toolbox LlamaIndex package\n",
        "!pip install toolbox-llamaindex --quiet\n",
        "\n",
        "# Install the llamaindex llm package\n",
        "# TODO(developer): replace this with another model if needed\n",
        "! pip install llama-index-llms-google-genai --quiet\n",
        "# ! pip install llama-index-llms-anthropic"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zjsq_xXice11"
      },
      "source": [
        "Create a LlamaIndex Hotel Agent which can Search, Book and Cancel hotels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EaBX4Dh6cU31",
        "outputId": "3d07e690-6102-4ed4-9751-fcf7fc869cdd"
      },
      "outputs": [],
      "source": [
        "import asyncio\n",
        "import os\n",
        "\n",
        "from llama_index.core.agent.workflow import AgentWorkflow\n",
        "\n",
        "from llama_index.core.workflow import Context\n",
        "\n",
        "# TODO(developer): replace this with another import if needed\n",
        "from llama_index.llms.google_genai import GoogleGenAI\n",
        "# from llama_index.llms.anthropic import Anthropic\n",
        "\n",
        "from toolbox_llamaindex import ToolboxClient\n",
        "\n",
        "prompt = \"\"\"\n",
        "  You're a helpful hotel assistant. You handle hotel searching, booking and\n",
        "  cancellations. When the user searches for a hotel, mention it's name, id,\n",
        "  location and price tier. Always mention hotel ids while performing any\n",
        "  searches. This is very important for any operations. For any bookings or\n",
        "  cancellations, please provide the appropriate confirmation. Be sure to\n",
        "  update checkin or checkout dates if mentioned by the user.\n",
        "  Don't ask for confirmations from the user.\n",
        "\"\"\"\n",
        "\n",
        "queries = [\n",
        "    \"Find hotels in Basel with Basel in it's name.\",\n",
        "    \"Can you book the Hilton Basel for me?\",\n",
        "    \"Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.\",\n",
        "    \"My check in dates would be from April 10, 2024 to April 19, 2024.\",\n",
        "]\n",
        "\n",
        "async def run_agent():\n",
        "    # Create an LLM to bind with the agent.\n",
        "    # TODO(developer): replace this with another model if needed\n",
        "    llm = GoogleGenAI(\n",
        "        model=\"gemini-2.0-flash-001\",\n",
        "        vertexai_config={\"project\": BIGQUERY_PROJECT, \"location\": \"us-central1\"},\n",
        "    )\n",
        "    # llm = GoogleGenAI(\n",
        "    #     api_key=os.getenv(\"GOOGLE_API_KEY\"),\n",
        "    #     model=\"gemini-2.0-flash-001\",\n",
        "    # )\n",
        "    # llm = Anthropic(\n",
        "    #   model=\"claude-3-7-sonnet-latest\",\n",
        "    #   api_key=os.getenv(\"ANTHROPIC_API_KEY\")\n",
        "    # )\n",
        "\n",
        "    # Load the tools from the Toolbox server\n",
        "    client = ToolboxClient(\"http://127.0.0.1:5000\")\n",
        "    tools = client.load_toolset()\n",
        "\n",
        "    # Create a LlamaIndex agent\n",
        "    agent = AgentWorkflow.from_tools_or_functions(\n",
        "        tools,\n",
        "        llm=llm,\n",
        "        system_prompt=prompt,\n",
        "    )\n",
        "\n",
        "    # Run the agent\n",
        "    ctx = Context(agent)\n",
        "    for query in queries:\n",
        "        response = await agent.run(user_msg=query, ctx=ctx)\n",
        "        print(f\"---- {query} ----\")\n",
        "        print(str(response))\n",
        "\n",
        "await run_agent()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yA3rAiELIds5"
      },
      "source": [
        "### Connect Using ADK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rphyouv2JtwX",
        "outputId": "63b26ec0-1880-4112-a90e-26b297ddd565"
      },
      "outputs": [],
      "source": [
        "!pip install -q google-adk\n",
        "!pip install -q toolbox-core"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wsc-bozAIXvo",
        "outputId": "6ecb10cd-2493-4bf4-e9ed-add9fed5c600"
      },
      "outputs": [],
      "source": [
        "# Create an ADK Hotel Agent which can Search, Book and Cancel hotels.\n",
        "from google.adk.agents import Agent\n",
        "from google.adk.runners import Runner\n",
        "from google.adk.sessions import InMemorySessionService\n",
        "from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService\n",
        "from google.genai import types\n",
        "from toolbox_core import ToolboxSyncClient\n",
        "\n",
        "import os\n",
        "\n",
        "os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = 'True'\n",
        "os.environ['GOOGLE_CLOUD_PROJECT'] = BIGQUERY_PROJECT\n",
        "os.environ['GOOGLE_CLOUD_LOCATION'] = 'us-central1'\n",
        "\n",
        "toolbox_client = ToolboxSyncClient(\"http://127.0.0.1:5000\")\n",
        "\n",
        "prompt = \"\"\"\n",
        "  You're a helpful hotel assistant. You handle hotel searching, booking and\n",
        "  cancellations. When the user searches for a hotel, mention it's name, id,\n",
        "  location and price tier. Always mention hotel ids while performing any\n",
        "  searches. This is very important for any operations. For any bookings or\n",
        "  cancellations, please provide the appropriate confirmation. Be sure to\n",
        "  update checkin or checkout dates if mentioned by the user.\n",
        "  Don't ask for confirmations from the user.\n",
        "\"\"\"\n",
        "\n",
        "root_agent = Agent(\n",
        "    model='gemini-2.0-flash-001',\n",
        "    name='hotel_agent',\n",
        "    description='A helpful AI assistant.',\n",
        "    instruction=prompt,\n",
        "    tools=toolbox_client.load_toolset(\"my-toolset\"),\n",
        ")\n",
        "\n",
        "session_service = InMemorySessionService()\n",
        "artifacts_service = InMemoryArtifactService()\n",
        "session = session_service.create_session(\n",
        "    state={}, app_name='hotel_agent', user_id='123'\n",
        ")\n",
        "runner = Runner(\n",
        "    app_name='hotel_agent',\n",
        "    agent=root_agent,\n",
        "    artifact_service=artifacts_service,\n",
        "    session_service=session_service,\n",
        ")\n",
        "\n",
        "queries = [\n",
        "    \"Find hotels in Basel with Basel in it's name.\",\n",
        "    \"Can you book the Hilton Basel for me?\",\n",
        "    \"Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.\",\n",
        "    \"My check in dates would be from April 10, 2024 to April 19, 2024.\",\n",
        "]\n",
        "\n",
        "for query in queries:\n",
        "    content = types.Content(role='user', parts=[types.Part(text=query)])\n",
        "    events = runner.run(session_id=session.id,\n",
        "                        user_id='123', new_message=content)\n",
        "\n",
        "    responses = (\n",
        "      part.text\n",
        "      for event in events\n",
        "      for part in event.content.parts\n",
        "      if part.text is not None\n",
        "    )\n",
        "\n",
        "    for text in responses:\n",
        "      print(text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kd-wF_Z9vVe3"
      },
      "source": [
        "### Observe the output\n",
        "\n",
        "You can see that the `Hyatt Regency Basel` has been booked for the correct dates."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 398
        },
        "id": "ZTW9bTUoqHis",
        "outputId": "0bd858b3-366e-4821-d570-3058e6bea673"
      },
      "outputs": [],
      "source": [
        "sql_select = f\"SELECT * FROM `{BIGQUERY_PROJECT}.{DATASET}.{TABLE_ID}`\"\n",
        "query_job = bqclient.query(sql_select)\n",
        "\n",
        "print(\"\\nQuery results:\")\n",
        "query_job.to_dataframe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wCRH0542KC51"
      },
      "source": [
        "### Clean-Up\n",
        "Conditionally delete BigQuery table and dataset in final session."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "w8W1N0rpz5Iq",
        "outputId": "0b6763d6-08c6-4c4d-e697-f6ae47fce3f2"
      },
      "outputs": [],
      "source": [
        "bqclient.delete_table(table_ref, not_found_ok=True)\n",
        "\n",
        "bqclient.get_dataset(dataset_ref)\n",
        "tables_in_dataset = list(bqclient.list_tables(dataset_ref))\n",
        "if not tables_in_dataset:\n",
        "    bqclient.delete_dataset(dataset_ref, delete_contents=False, not_found_ok=True)\n",
        "    print(f\"Dataset '{DATASET}' deleted.\")\n",
        "else:\n",
        "    print(f\"Dataset '{DATASET}' is not empty. Skipping dataset deletion.\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
