MCP Toolbox for Databases

{ "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": [ "[](https://colab.research.google.com/github/googleapis/genai-toolbox/blob/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": {}, "source": [ "## Optional: Enable Vertex AI API for Google Cloud\n", "\n", "If you're using a model hosted on **Vertex AI**, run the following command to enable the API:\n", "\n", "```bash\n", "!gcloud services enable aiplatform.googleapis.com\n" ] }, { "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": [ "version = \"0.17.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": "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=Fv2-uT4mvYtp)\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": "QqRlWqvYNKSo" }, "source": [ "### Connect Using ADK" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "dhQTKlpVNKSo" }, "outputs": [], "source": [ "! pip install toolbox-core --quiet\n", "! pip install google-adk --quiet" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tSLO_0vKNKSo" }, "outputs": [], "source": [ "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", "# TODO(developer): replace this with your Google API key\n", "os.environ['GOOGLE_API_KEY'] = \"<GOOGLE_API_KEY>\"\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 = await 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-2.0-flash-001\", project=project_id)\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 = 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-2.0-flash-001\",\n", " vertexai_config={\"project\": project_id, \"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 = 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": "Fv2-uT4mvYtp" }, "source": [ "### Connect Using Google GenAI" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "mHSvk5_AvYtu" }, "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": "sO_7FGSYvYtu" }, "source": [ "Create a Google GenAI Application which can Search, Book and Cancel hotels." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "-NVVBiLnvYtu" }, "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-001\",\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": "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 }