MCP Toolbox for Databases

{ "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.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": "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 }