docs/en/getting-started/colab_quickstart.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KWS7OXcEJptT"
      },
      "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": "5sZ7_HCYJm4y"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/googleapis/genai-toolbox/tree/main/docs/en/getting-started/colab_quickstart.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DMLivT-MIcmV"
      },
      "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, PostgreSQL, 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 `PostgreSQL database`.\n",
        "- Launch a Toolbox server.\n",
        "- Connect to Toolbox and develop a sample `Hotel Booking` application.\n",
        "\n",
        "Here is the simplified flow of a Toolbox Application:\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/toolbox_flow.png\" alt=\"Toolbox Flow\"/>\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nBWNgqEWIoSG"
      },
      "source": [
        "## Step 1: Set up your database\n",
        "\n",
        "In this section, we will\n",
        "1. Create a database.\n",
        "1. Create a user to access the database.\n",
        "1. Insert example data into the database."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EeFDe-j4HUmn"
      },
      "outputs": [],
      "source": [
        "# Install postgresql to run a DB server on colab\n",
        "%%shell\n",
        "\n",
        "sudo apt-get -y -qq update > /dev/null 2>&1\n",
        "sudo apt-get -y -qq install postgresql > /dev/null 2>&1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5D9t-ZB1Dre-"
      },
      "outputs": [],
      "source": [
        "# Start the postgresql server.\n",
        "!sudo service postgresql start"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qJJbo5T7sjyd"
      },
      "outputs": [],
      "source": [
        "# Check that postgres is running\n",
        "!sudo lsof -i :5432"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zTtKdvbwAag3"
      },
      "outputs": [],
      "source": [
        "# Create a dedicated database and a user to access our DB securely\n",
        "%%shell\n",
        "\n",
        "sudo -u postgres psql << EOF\n",
        "CREATE USER toolbox_user WITH PASSWORD 'my-password';\n",
        "CREATE DATABASE toolbox_db;\n",
        "GRANT ALL PRIVILEGES ON DATABASE toolbox_db TO toolbox_user;\n",
        "ALTER DATABASE toolbox_db OWNER TO toolbox_user;\n",
        "EOF"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zVx18ijjrWSO"
      },
      "source": [
        "> **Tip:** For a real application, it’s best to follow the principle of least permission and only grant the privileges your application needs.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6t_nLJIHCRgy"
      },
      "outputs": [],
      "source": [
        "# Connect to the database with the new user and create a hotels table.\n",
        "%%shell\n",
        "\n",
        "export PGPASSWORD=my-password\n",
        "psql -h 127.0.0.1 -U toolbox_user -d toolbox_db --no-password << EOF\n",
        "CREATE TABLE hotels(\n",
        "   id            INTEGER NOT NULL PRIMARY KEY,\n",
        "   name          VARCHAR NOT NULL,\n",
        "   location      VARCHAR NOT NULL,\n",
        "   price_tier    VARCHAR NOT NULL,\n",
        "   checkin_date  DATE    NOT NULL,\n",
        "   checkout_date DATE    NOT NULL,\n",
        "   booked        BIT     NOT NULL\n",
        ");\n",
        "INSERT INTO hotels(id, name, location, price_tier, checkin_date, checkout_date, booked)\n",
        "VALUES\n",
        "  (1, 'Hilton Basel', 'Basel', 'Luxury', '2024-04-22', '2024-04-20', B'0'),\n",
        "  (2, 'Marriott Zurich', 'Zurich', 'Upscale', '2024-04-14', '2024-04-21', B'0'),\n",
        "  (3, 'Hyatt Regency Basel', 'Basel', 'Upper Upscale', '2024-04-02', '2024-04-20', B'0'),\n",
        "  (4, 'Radisson Blu Lucerne', 'Lucerne', 'Midscale', '2024-04-24', '2024-04-05', B'0'),\n",
        "  (5, 'Best Western Bern', 'Bern', 'Upper Midscale', '2024-04-23', '2024-04-01', B'0'),\n",
        "  (6, 'InterContinental Geneva', 'Geneva', 'Luxury', '2024-04-23', '2024-04-28', B'0'),\n",
        "  (7, 'Sheraton Zurich', 'Zurich', 'Upper Upscale', '2024-04-27', '2024-04-02', B'0'),\n",
        "  (8, 'Holiday Inn Basel', 'Basel', 'Upper Midscale', '2024-04-24', '2024-04-09', B'0'),\n",
        "  (9, 'Courtyard Zurich', 'Zurich', 'Upscale', '2024-04-03', '2024-04-13', B'0'),\n",
        "  (10, 'Comfort Inn Bern', 'Bern', 'Midscale', '2024-04-04', '2024-04-16', B'0');\n",
        "SELECT * from hotels;\n",
        "EOF"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rw5fBXxpJ3-Q"
      },
      "outputs": [],
      "source": [
        "# Check that database is running\n",
        "!sudo lsof -i :5432"
      ]
    },
    {
      "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": {
        "id": "lbsQ1Aa-IszB"
      },
      "outputs": [],
      "source": [
        "! curl -O https://storage.googleapis.com/genai-toolbox/v0.2.0/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": "4HQwVY5Pu1Xi"
      },
      "source": [
        "> Note: To include a literal dollar sign (e.g., $1) as part of your SQL statement within the Python string for tools.yml, you must escape both the backslash and the dollar sign. Use \\\\\\$1 in Python to output \\$1 in the tools.yml file.\n",
        "\n",
        "> Note: You can also set up Colab secrets to store any sensitive information like passwords. You can easily add secrets through the left panel:\n",
        "\n",
        "<img src=\"https://services.google.com/fh/files/misc/colab_secret.png\" alt=\"Colab Secrets\" width=\"400\"/>\n"
      ]
    },
    {
      "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 database.\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-pg-source:\n",
        "    kind: postgres\n",
        "    host: 127.0.0.1\n",
        "    port: 5432\n",
        "    database: toolbox_db\n",
        "    user: toolbox_user\n",
        "    password: my-password\n",
        "tools:\n",
        "  search-hotels-by-name:\n",
        "    kind: postgres-sql\n",
        "    source: my-pg-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 hotels WHERE name ILIKE '%' || \\$1 || '%';\n",
        "  search-hotels-by-location:\n",
        "    kind: postgres-sql\n",
        "    source: my-pg-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 hotels WHERE location ILIKE '%' || \\$1 || '%';\n",
        "  book-hotel:\n",
        "    kind: postgres-sql\n",
        "    source: my-pg-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: string\n",
        "        description: The ID of the hotel to book.\n",
        "    statement: UPDATE hotels SET booked = B'1' WHERE id = \\$1;\n",
        "  update-hotel:\n",
        "    kind: postgres-sql\n",
        "    source: my-pg-source\n",
        "    description: >-\n",
        "      Update a hotel's check-in and check-out dates by its ID. Returns a message\n",
        "      indicating  whether the hotel was successfully updated or not.\n",
        "    parameters:\n",
        "      - name: hotel_id\n",
        "        type: string\n",
        "        description: The ID of the hotel to update.\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",
        "    statement: >-\n",
        "      UPDATE hotels SET checkin_date = CAST(\\$2 as date), checkout_date = CAST(\\$3\n",
        "      as date) WHERE id = \\$1;\n",
        "  cancel-hotel:\n",
        "    kind: postgres-sql\n",
        "    source: my-pg-source\n",
        "    description: Cancel a hotel by its ID.\n",
        "    parameters:\n",
        "      - name: hotel_id\n",
        "        type: string\n",
        "        description: The ID of the hotel to cancel.\n",
        "    statement: UPDATE hotels SET booked = B'0' WHERE id = \\$1;\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": [
        "# Write the file content into the tools file.\n",
        "! echo \"{file_content}\" > \"{tools_file_name}\""
      ]
    },
    {
      "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": {
        "id": "1PJpKOBieKOV"
      },
      "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": "markdown",
      "metadata": {
        "id": "yfg-u9Y4Mu_a"
      },
      "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": "ky64V76MMttC"
      },
      "outputs": [],
      "source": [
        "# Run this and allow access through the pop-up\n",
        "from google.colab import auth\n",
        "\n",
        "auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u0Jc-0YNdhQd"
      },
      "outputs": [],
      "source": [
        "# @markdown Please fill in the value below with your GCP project ID and then run the cell.\n",
        "\n",
        "# Please fill in these values.\n",
        "project_id = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Quick input validations.\n",
        "assert project_id, \"⚠️ Please provide a Google Cloud project ID\"\n",
        "\n",
        "# Configure gcloud.\n",
        "!gcloud config set project {project_id}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J46eLkFbNhWq"
      },
      "source": [
        "> You can either use LangGraph or LlamaIndex to develop a Toolbox based\n",
        "> application. Run one of the sections below\n",
        "> - [Connect using Google GenAI](#scrollTo=Rwgv1LDdNKSn)\n",
        "> - [Connect using ADK](#scrollTo=QqRlWqvYNKSo)\n",
        "> - [Connect Using LangGraph](#scrollTo=pbapNMhhL33S)\n",
        "> - [Connect using LlamaIndex](#scrollTo=04iysrm_L_7v)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rwgv1LDdNKSn"
      },
      "source": [
        "### Connect Using Google GenAI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HY23RMk4NKSn"
      },
      "outputs": [],
      "source": [
        "# Install the Toolbox Core package\n",
        "!pip install toolbox-core --quiet\n",
        "\n",
        "# Install the Google GenAI package\n",
        "!pip install google-genai --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9F1u566sNKSn"
      },
      "source": [
        "Create a Google GenAI Application which can Search, Book and Cancel hotels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LAuBIOXvNKSn"
      },
      "outputs": [],
      "source": [
        "import asyncio\n",
        "\n",
        "from google import genai\n",
        "from google.genai.types import (\n",
        "    Content,\n",
        "    FunctionDeclaration,\n",
        "    GenerateContentConfig,\n",
        "    Part,\n",
        "    Tool,\n",
        ")\n",
        "\n",
        "from toolbox_core 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",
        "    \"Please book the hotel Hilton Basel for me.\",\n",
        "    \"This is too expensive. Please cancel it.\",\n",
        "    \"Please book Hyatt Regency for me\",\n",
        "    \"My check in dates for my booking would be from April 10, 2024 to April 19, 2024.\",\n",
        "]\n",
        "\n",
        "\n",
        "async def run_application():\n",
        "    toolbox_client = ToolboxClient(\"http://127.0.0.1:5000\")\n",
        "\n",
        "    # The toolbox_tools list contains Python callables (functions/methods) designed for LLM tool-use\n",
        "    # integration. While this example uses Google's genai client, these callables can be adapted for\n",
        "    # various function-calling or agent frameworks. For easier integration with supported frameworks\n",
        "    # (https://github.com/googleapis/mcp-toolbox-python-sdk/tree/main/packages), use the\n",
        "    # provided wrapper packages, which handle framework-specific boilerplate.\n",
        "    toolbox_tools = await toolbox_client.load_toolset(\"my-toolset\")\n",
        "    genai_client = genai.Client(\n",
        "        vertexai=True, project=project_id, location=\"us-central1\"\n",
        "    )\n",
        "\n",
        "    genai_tools = [\n",
        "        Tool(\n",
        "            function_declarations=[\n",
        "                FunctionDeclaration.from_callable_with_api_option(callable=tool)\n",
        "            ]\n",
        "        )\n",
        "        for tool in toolbox_tools\n",
        "    ]\n",
        "    history = []\n",
        "    for query in queries:\n",
        "        user_prompt_content = Content(\n",
        "            role=\"user\",\n",
        "            parts=[Part.from_text(text=query)],\n",
        "        )\n",
        "        history.append(user_prompt_content)\n",
        "\n",
        "        response = genai_client.models.generate_content(\n",
        "            model=\"gemini-2.0-flash\",\n",
        "            contents=history,\n",
        "            config=GenerateContentConfig(\n",
        "                system_instruction=prompt,\n",
        "                tools=genai_tools,\n",
        "            ),\n",
        "        )\n",
        "        history.append(response.candidates[0].content)\n",
        "        function_response_parts = []\n",
        "        for function_call in response.function_calls:\n",
        "            fn_name = function_call.name\n",
        "            # The tools are sorted alphabetically\n",
        "            if fn_name == \"search-hotels-by-name\":\n",
        "                function_result = await toolbox_tools[3](**function_call.args)\n",
        "            elif fn_name == \"search-hotels-by-location\":\n",
        "                function_result = await toolbox_tools[2](**function_call.args)\n",
        "            elif fn_name == \"book-hotel\":\n",
        "                function_result = await toolbox_tools[0](**function_call.args)\n",
        "            elif fn_name == \"update-hotel\":\n",
        "                function_result = await toolbox_tools[4](**function_call.args)\n",
        "            elif fn_name == \"cancel-hotel\":\n",
        "                function_result = await toolbox_tools[1](**function_call.args)\n",
        "            else:\n",
        "                raise ValueError(\"Function name not present.\")\n",
        "            function_response = {\"result\": function_result}\n",
        "            function_response_part = Part.from_function_response(\n",
        "                name=function_call.name,\n",
        "                response=function_response,\n",
        "            )\n",
        "            function_response_parts.append(function_response_part)\n",
        "\n",
        "        if function_response_parts:\n",
        "            tool_response_content = Content(role=\"tool\", parts=function_response_parts)\n",
        "            history.append(tool_response_content)\n",
        "\n",
        "        response2 = genai_client.models.generate_content(\n",
        "            model=\"gemini-2.0-flash-001\",\n",
        "            contents=history,\n",
        "            config=GenerateContentConfig(\n",
        "                tools=genai_tools,\n",
        "            ),\n",
        "        )\n",
        "        final_model_response_content = response2.candidates[0].content\n",
        "        history.append(final_model_response_content)\n",
        "        print(response2.text)\n",
        "\n",
        "\n",
        "asyncio.run(run_application())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QqRlWqvYNKSo"
      },
      "source": [
        "### Connect Using ADK"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dhQTKlpVNKSo"
      },
      "outputs": [],
      "source": [
        "! pip install toolbox-langchain --quiet\n",
        "! pip install google-adk langchain --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tSLO_0vKNKSo"
      },
      "outputs": [],
      "source": [
        "from google.adk.agents import Agent\n",
        "from google.adk.tools.toolbox_tool import ToolboxTool\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",
        "\n",
        "import os\n",
        "# TODO(developer): replace this with your Google API key\n",
        "os.environ['GOOGLE_API_KEY'] = \"<GOOGLE_API_KEY>\"\n",
        "\n",
        "toolbox_tools = ToolboxTool(\"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',\n",
        "    name='hotel_agent',\n",
        "    description='A helpful AI assistant.',\n",
        "    instruction=prompt,\n",
        "    tools=toolbox_tools.get_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": "pbapNMhhL33S"
      },
      "source": [
        "### Connect Using LangGraph"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uraBx8mbMXnV"
      },
      "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": {
        "id": "Br3ucM46M9uc"
      },
      "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",
        "async def run_application():\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-1.5-pro\", project=project_id)\n",
        "  # model = ChatGoogleGenerativeAI(model=\"gemini-1.5-pro\")\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 = await client.aload_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)\n",
        "\n",
        "await run_application()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "04iysrm_L_7v"
      },
      "source": [
        "### Connect using LlamaIndex"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6b6Loh8SJ_iA"
      },
      "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": {
        "id": "EaBX4Dh6cU31"
      },
      "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_application():\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-1.5-pro\",\n",
        "        vertexai_config={\"project\": project_id, \"location\": \"us-central1\"},\n",
        "    )\n",
        "    # llm = GoogleGenAI(\n",
        "    #     api_key=os.getenv(\"GOOGLE_API_KEY\"),\n",
        "    #     model=\"gemini-1.5-pro\",\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 = await client.aload_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_application()"
      ]
    },
    {
      "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": {
        "id": "ZTW9bTUoqHis"
      },
      "outputs": [],
      "source": [
        "%%shell\n",
        "\n",
        "export PGPASSWORD=my-password\n",
        "psql -h 127.0.0.1 -U toolbox_user -d toolbox_db --no-password << EOF\n",
        "SELECT * from hotels;\n",
        "EOF"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qV36Do-Bub12"
      },
      "source": [
        "## Optional: Cleanup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yatf9YoGclV9"
      },
      "source": [
        "Executing this will terminate the processes running on the database and Toolbox ports.\n",
        "\n",
        "This is necessary before re-running the startup cells for these services to prevent `port already in use` errors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WC8doSdzJkDE"
      },
      "outputs": [],
      "source": [
        "!lsof -t -i :5432 | xargs kill -9"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D09SAmLzufwO"
      },
      "outputs": [],
      "source": [
        "# Verify that the database process is killed\n",
        "!sudo lsof -i :5432"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Rwgv1LDdNKSn",
        "pbapNMhhL33S",
        "04iysrm_L_7v"
      ],
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}