amazon-bedrock-agentcore-samples/01-tutorials/02-AgentCore-gateway/03-search-tools/01-gateway-search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0f80631b",
   "metadata": {},
   "source": [
    "# Amazon Bedrock AgentCore Gateway - Semantic search tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b6a992b",
   "metadata": {},
   "source": [
    "### Tutorial Details\n",
    "\n",
    "\n",
    "| Information         | Details                                                                          |\n",
    "|:--------------------|:---------------------------------------------------------------------------------|\n",
    "| Tutorial type       | Conversational                                                                   |\n",
    "| Agent type          | Single                                                                           |\n",
    "| AgentCore services  | AgentCore Gateway, AgentCore Identity                                            |\n",
    "| Agentic Framework   | Strands Agents                                                                   |\n",
    "| LLM model           | Anthropic Claude Sonnet 3.7                                                        |\n",
    "| Tutorial components | Creating and using Lambda-backed AgentCore Gateway from Strands Agent            |\n",
    "| Tutorial vertical   | Cross-vertical                                                                   |\n",
    "| Example complexity  | Easy                                                                             |\n",
    "| SDK used            | Amazon BedrockAgentCore Python SDK and boto3  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61187ab7",
   "metadata": {},
   "source": [
    "### Tutorial Architecture\n",
    "Amazon Bedrock AgentCore Gateway provides unified connectivity between agents and the tools and resources they need to interact with. Gateway plays multiple roles in this connectivity layer:\n",
    "\n",
    "1. **Security Guard**: Gateway manages OAuth authorization to ensure only valid users / agents access tools / resources.\n",
    "2. **Translator**: Gateway translates agent requests made using popular protocols like the Model Context Protocol (MCP) into API requests and Lambda invocations. This means developers don’t need to host servers, manage protocol integration, version support, version patching, etc.\n",
    "3. **Composer**: Gateway enables developers to seamlessly combine multiple APIs, functions, and tools into a single MCP  endpoint that an agent can use.\n",
    "4. **Keychain**: Gateway handles the injection of the right credentials to use with the right tool, ensuring that agents can seamlessly leverage tools that require different sets of credentials. \n",
    "5. **Researcher**: Gateway enables agents to search across all of their tools to find only the ones that are best for a given context or question. This allows agents to make use of 1000s of tools instead of just a handful. It also minimizes the set of tools that need to be provided in an agent’s LLM prompt, reducing latency and cost. \n",
    "6. **Infrastructure Manager**: Gateway is completely serverless, and comes with built-in observability and auditing, alleviating the need for developers to manage additional infrastructure to integrate their agents and tools. \n",
    "\n",
    "![How does it work](images/gw-arch-overview.png)\n",
    "\n",
    "### Tutorial Key Features\n",
    "\n",
    "* Creating Amazon Bedrock AgentCore Gateways with AWS Lambda-backed targets\n",
    "* Using AgentCore Gateway semantic search \n",
    "* Using Strands Agents to show how AgentCore Gateway search improves latency"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa1feb54",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "To execute this tutorial you will need:\n",
    "* Python 3.10+\n",
    "* AWS credentials\n",
    "* Amazon Bedrock AgentCore SDK\n",
    "* Strands Agents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb01ea0d",
   "metadata": {},
   "source": [
    "## AgentCore Gateway helps solve the challenge of MCP servers that have large numbers of tools\n",
    "In a typical enterprise setting, agent builders encounter MCP servers that have hundreds or even thousands\n",
    "of MCP tools. This volume of tools poses challenges for AI agents, including poor tool selection accuracy, \n",
    "increased cost, and higher latency driven by higher token usage from excessive tool metadata.\n",
    "This can happen when connecting your agents to third party services (e.g., Zendesk, Salesforce,\n",
    "Slack, JIRA, ...), or to existing enterprise REST services. \n",
    "AgentCore Gateway provides a built in semantic search across tools, \n",
    "which improves agent latency, cost, and accuracy, while still giving those agents the tools they need. \n",
    "Depending on your use case, LLM model, and agent framework, you can see up to 3x better latency by keeping\n",
    "an agent focused on relevant tools versus providing the full set of hundreds of tools from a typical MCP Server.\n",
    "\n",
    "![How does it work](images/gateway_tool_search.png)\n",
    "\n",
    "## What you will learn in this notebook\n",
    "In this notebook, we provide a tutorial for AgentCore Gateway search. By the end of this step-by-step tutorial, you\n",
    "will understand:\n",
    "\n",
    "- How to use AgentCore Gateway's built-in search tool to quickly find relevant tools \n",
    "- How to integrate tool search results into Strands Agents for improved latency and reduced cost\n",
    "\n",
    "## Overview of the notebook structure\n",
    "The notebook is structured with the following sections:\n",
    "\n",
    "1. Understanding fundamentals of AgentCore Gateway Search\n",
    "2. Preparing the notebook environment\n",
    "3. Setting up a Gateway that has hundreds of tools\n",
    "4. Searching for tools from a Gateway\n",
    "5. Using Strands Agents with an MCP server that has many tools\n",
    "6. Adding tool search results to a Strands Agent\n",
    "7. Showing 3x latency improvement by using tool search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06b7852f",
   "metadata": {},
   "source": [
    "# Understanding fundamentals of AgentCore Gateway Search\n",
    "\n",
    "When you create an AgentCore Gateway, you have the option to indicate that you want Search enabled.\n",
    "For Gateways with search enabled, three things happen:\n",
    "\n",
    "1. **Vector store is created**. The Gateway service automatically creates a serverless fully-managed vector store for your new Gateway. This enables a full semantic search across your Gateway tools. \n",
    "3. **Vector store is populated**. As you add Gateway Targets to your Gateway, the service automatically uses embeddings behind the scenes to populate the vector store based on the tools from the new Target. The tool metadata comes from the JSON defintions of your tools or the OpenAPI Schema specification for your REST services targets.\n",
    "2. **Search tool (MCP based) is provided**. In addition to all of your user defined tools (from AWS Lambda targets or REST services), the Gateway gets one additional MCP tool that provides semantic search. It is named `x-amz-bedrock-agentcore-search`. The prefix ensures there are no name clashes with your user-defined tools. We may add more tools like that in the future as well. The search tool has a single argument called `query`. When the search tool is invoked, the Gateway service performs a semantic search using that query, matching it against available tool metadata (names, descriptions, input and output schema), and returns the most relevant tools in descending order of relevance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b39c756c",
   "metadata": {},
   "source": [
    "# Preparing the notebook environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2045271e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --force-reinstall -U -r requirements.txt --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dec9b4b",
   "metadata": {},
   "source": [
    "Import all the required Python libraries, and load environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10100b26",
   "metadata": {},
   "outputs": [],
   "source": [
    "from strands import Agent\n",
    "from strands.models import BedrockModel\n",
    "from strands.handlers import null_callback_handler\n",
    "\n",
    "from strands.tools.mcp.mcp_client import MCPClient, MCPAgentTool\n",
    "\n",
    "from mcp.client.streamable_http import streamablehttp_client\n",
    "from mcp.types import Tool as MCPTool\n",
    "\n",
    "import logging\n",
    "import time\n",
    "import json\n",
    "import boto3\n",
    "import requests\n",
    "import utils\n",
    "\n",
    "GATEWAY_NAME = \"gateway-search-tutorial\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e76ded57",
   "metadata": {},
   "source": [
    "Set up a logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea5ab3da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure the root strands logger\n",
    "logging.getLogger(\"strands\").setLevel(logging.ERROR)  # INFO) #DEBUG) #\n",
    "\n",
    "# Add a handler to see the logs\n",
    "logging.basicConfig(\n",
    "    format=\"%(levelname)s | %(name)s | %(message)s\", handlers=[logging.StreamHandler()]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dc97c60",
   "metadata": {},
   "source": [
    "Check our boto3 version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9258624c",
   "metadata": {},
   "outputs": [],
   "source": [
    "boto3.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59719b74",
   "metadata": {},
   "source": [
    "Get our boto3 client for the AgentCore control plane API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2eafe7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = boto3.Session()\n",
    "agentcore_client = session.client(\n",
    "    \"bedrock-agentcore-control\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb9fec29",
   "metadata": {},
   "source": [
    "# Setting up a Gateway that has hundreds of tools\n",
    "\n",
    "AgentCore Gateway provides a secure and scalable way to expose a curated set of existing APIs\n",
    "as MCP tools for your agents. In a production setting, your Gateway resources would be created\n",
    "using infrastructure as code with tools like CloudFormation, CDK, or Terraform. In this tutorial,\n",
    "we use the boto3 control plane APIs directly, so that you can understand the resources and APIs more effectively.\n",
    "This will help you more easily get started building and using your own gateways, and making more powerful\n",
    "and secure agents.\n",
    "\n",
    "At a high level, the steps for setting up your Gateway are:\n",
    "\n",
    "1. Define what identity providers and credential providers you are using for inbound (agents calling Gateways) and outbound (Gateways calling tools) security.\n",
    "2. Create the Gateway using `create_gateway`.\n",
    "3. Add Gateway Targets using `create_gateway_target`, to expose MCP tools that will be implemented in AWS Lambda or in existing RESTful services.\n",
    "\n",
    "In this tutorial, we will use Amazon Cognito as the identity provider (IdP), AWS Lambda functions as targets, and AWS IAM for outbound authentication. The same concepts demonstrated in this tutorial still apply when using other IdP's or other target types."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d896e613",
   "metadata": {},
   "source": [
    "### Creating Amazon Cognito resources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb02989f",
   "metadata": {},
   "source": [
    "In this tutorial, we assume you have already created the following resources and have set up corresponding environment variables:\n",
    "\n",
    "- IAM role for AWS Lambda execution (`gateway_lambda_iam_role`)\n",
    "- AWS Lambda function for your simple math tools (`calc_lambda_arn`)\n",
    "- AWS Lambda function for your restaurant reservation tool (`restaurant_lambda_arn`)\n",
    "- Amazon Cognito user pools, giving you a client id (`cognito_client_id`) and a discovery URL (`cognito_discovery_url`)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9cb9eb5",
   "metadata": {},
   "source": [
    "Lets take a look at the JSON tool metadata for the restaurant API. Note that if we were integrating with existing REST services, the API specs would be provided using OpenAPI Schema instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e953add5",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./restaurant/restaurant-api.json\") as f:\n",
    "    data = json.load(f)\n",
    "print(json.dumps(data, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0758917",
   "metadata": {},
   "source": [
    "Here are the simple calculator APIs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef1a47e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./calc/calc-api.json\") as f:\n",
    "    data = json.load(f)[0:3]\n",
    "print(json.dumps(data, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "414b7a8a",
   "metadata": {},
   "source": [
    "Here is the AWS Lambda function implementation for the calculator tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48160295",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, Code\n",
    "\n",
    "with open(\"./calc/lambda_function_code.py\", \"r\") as f:\n",
    "    code_content = f.read()\n",
    "display(Code(code_content, language=\"python\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b25b7e35",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"./restaurant/lambda_function_code.py\", \"r\") as f:\n",
    "    code_content = f.read()\n",
    "display(Code(code_content, language=\"python\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f400f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Create a sample AWS Lambda function that you want to convert into MCP tools\n",
    "calc_lambda_resp = utils.create_gateway_lambda(\n",
    "    \"calc/lambda_function_code.zip\", lambda_function_name=\"calc_lambda_gateway\"\n",
    ")\n",
    "\n",
    "if calc_lambda_resp is not None:\n",
    "    if calc_lambda_resp[\"exit_code\"] == 0:\n",
    "        print(\n",
    "            \"Lambda function created with ARN: \",\n",
    "            calc_lambda_resp[\"lambda_function_arn\"],\n",
    "        )\n",
    "    else:\n",
    "        print(\n",
    "            \"Lambda function creation failed with message: \",\n",
    "            calc_lambda_resp[\"lambda_function_arn\"],\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "176e24bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "calc_lambda_resp[\"lambda_function_arn\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2db00b43",
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Create a sample AWS Lambda function that you want to convert into MCP tools\n",
    "restaurant_lambda_resp = utils.create_gateway_lambda(\n",
    "    \"restaurant/lambda_function_code.zip\",\n",
    "    lambda_function_name=\"restaurant_lambda_gateway\",\n",
    ")\n",
    "\n",
    "if restaurant_lambda_resp is not None:\n",
    "    if restaurant_lambda_resp[\"exit_code\"] == 0:\n",
    "        print(\n",
    "            \"Lambda function created with ARN: \",\n",
    "            restaurant_lambda_resp[\"lambda_function_arn\"],\n",
    "        )\n",
    "    else:\n",
    "        print(\n",
    "            \"Lambda function creation failed with message: \",\n",
    "            restaurant_lambda_resp[\"lambda_function_arn\"],\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc46a911",
   "metadata": {},
   "outputs": [],
   "source": [
    "restaurant_lambda_resp[\"lambda_function_arn\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a929d763",
   "metadata": {},
   "outputs": [],
   "source": [
    "cognito_response = utils.setup_cognito_user_pool()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1b3ba4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "bearer_token = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d2571fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "gateway_role_arn = utils.create_gateway_iam_role(\n",
    "    lambda_arns=[\n",
    "        calc_lambda_resp[\"lambda_function_arn\"],\n",
    "        restaurant_lambda_resp[\"lambda_function_arn\"],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9de9ef7",
   "metadata": {},
   "source": [
    "#### Let's create a few helper functions for using the control plane APIs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6a3cb63",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_apispec(json_file_path):\n",
    "    try:\n",
    "        # read json file and return contents as string\n",
    "        with open(json_file_path, \"r\") as file:\n",
    "            # Parse JSON to Python object\n",
    "            api_spec = json.load(file)\n",
    "            return api_spec\n",
    "\n",
    "    except FileNotFoundError:\n",
    "        return f\"Error: File {json_file_path} not found\"\n",
    "    except Exception as e:\n",
    "        return f\"An unexpected error occurred: {str(e)}\"\n",
    "\n",
    "\n",
    "def list_gateways():\n",
    "    response = agentcore_client.list_gateways()\n",
    "    print(json.dumps(response, indent=2, default=str))\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60d8aa76",
   "metadata": {},
   "source": [
    "#### Create Gateway helper function\n",
    "Here is a helper function for creating an AgentCore Gateway given a name and\n",
    "description. It uses Amazon Cognito as its IdP, and pulls the allowed client ID\n",
    "and the discovery URL from environment variables, as those were already defined.\n",
    "It also defaults to enabling semantic search on the resulting Gateway, and uses\n",
    "a predefined IAM role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffac7f83",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_gateway(gateway_name, gateway_desc):\n",
    "    # Use Cognito for Inbound OAuth to our Gateway\n",
    "    auth_config = {\n",
    "        \"customJWTAuthorizer\": {\n",
    "            \"allowedClients\": [cognito_response[\"client_id\"]],\n",
    "            \"discoveryUrl\": cognito_response[\"discovery_url\"],\n",
    "        }\n",
    "    }\n",
    "    # Enable semantic search of tools\n",
    "    search_config = {\n",
    "        \"mcp\": {\"searchType\": \"SEMANTIC\", \"supportedVersions\": [\"2025-03-26\"]}\n",
    "    }\n",
    "    # Create the gateway\n",
    "    response = agentcore_client.create_gateway(\n",
    "        name=gateway_name,\n",
    "        roleArn=gateway_role_arn,\n",
    "        authorizerType=\"CUSTOM_JWT\",\n",
    "        description=gateway_desc,\n",
    "        protocolType=\"MCP\",\n",
    "        authorizerConfiguration=auth_config,\n",
    "        protocolConfiguration=search_config,\n",
    "    )\n",
    "    # print(json.dumps(response, indent=2, default=str))\n",
    "    return response[\"gatewayId\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d45dd92",
   "metadata": {},
   "source": [
    "#### Create Gateway Target helper function\n",
    "This function creates a new AWS Lambda target on an existing Gateway.\n",
    "Simply provide the gateway ID, the name and description of the new target,\n",
    "the ARN of the existing AWS Lambda function, and the JSON schema describing\n",
    "the interfaces to the tools you want to expose from the gateway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d03d7c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_gatewaytarget(gateway_id, target_name, target_descr, lambda_arn, api_spec):\n",
    "    # Add a Lambda target to the gateway\n",
    "    response = agentcore_client.create_gateway_target(\n",
    "        gatewayIdentifier=gateway_id,\n",
    "        name=target_name,\n",
    "        description=target_descr,\n",
    "        targetConfiguration={\n",
    "            \"mcp\": {\n",
    "                \"lambda\": {\n",
    "                    \"lambdaArn\": lambda_arn,\n",
    "                    \"toolSchema\": {\"inlinePayload\": api_spec},\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        # Use IAM as credential provider\n",
    "        credentialProviderConfigurations=[\n",
    "            {\"credentialProviderType\": \"GATEWAY_IAM_ROLE\"}\n",
    "        ],\n",
    "    )\n",
    "    return response[\"targetId\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2084feec",
   "metadata": {},
   "source": [
    "### Creating your first AgentCore Gateway\n",
    "Before we setup your first Gateway, let's take a quick look at how \n",
    "Gateway provides security, both for inbound requests to use MCP tools,\n",
    "and outbound access from the Gateway to tools and resources.\n",
    "\n",
    "![How does it work](images/gateway_secure_access.png)\n",
    "\n",
    "Now let's create the gateway for this tutorial, providing a name and a description."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "175d1ef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Create gateway with name: {GATEWAY_NAME}\")\n",
    "gatewayId = create_gateway(\n",
    "    gateway_name=GATEWAY_NAME, gateway_desc=\"AgentCore Gateway Tutorial\"\n",
    ")\n",
    "print(f\"Gateway created with id: {gatewayId}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c878491f",
   "metadata": {},
   "source": [
    "### Adding AgentCore Gateway Targets\n",
    "In this tutorial, we assume you already have installed a pair of Lambda functions, one for doing simple\n",
    "math calculations, and another that simulates creating a restaurant reservation. We'll add a Gateway Target\n",
    "for each of these functions.\n",
    "\n",
    "Once we've added those targets, we'll add additional targets to simply drive higher MCP tool counts that\n",
    "help us demonstrate the power of AgentCore Gateway search."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4d24edb",
   "metadata": {},
   "source": [
    "Now that the gateway is created, let's add a target for making restaurant reservations\n",
    "via a Lambda function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58d04ee4",
   "metadata": {},
   "outputs": [],
   "source": [
    "restaurant_api_spec = read_apispec(\"./restaurant/restaurant-api.json\")\n",
    "restaurant_lambda_arn = restaurant_lambda_resp[\"lambda_function_arn\"]\n",
    "print(f\"Restaurant Lambda ARN: {restaurant_lambda_arn}\")\n",
    "\n",
    "restaurantTargetId = create_gatewaytarget(\n",
    "    gateway_id=gatewayId,\n",
    "    lambda_arn=restaurant_lambda_arn,\n",
    "    target_name=\"FoodTools\",\n",
    "    target_descr=\"Restaurant Tools\",\n",
    "    api_spec=restaurant_api_spec,\n",
    ")\n",
    "print(f\"RestaurantTarget created with id: {restaurantTargetId} on gateway: {gatewayId}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4da14900",
   "metadata": {},
   "source": [
    "Here we'll add a second target, this time with a Lambda that implements 4 basic tools (add, subtract,\n",
    "multiply, divide), and a set of 75 generated tool definitions for investment management (trading, credit research, quantitative analysis, portfolio management). The investment management tool definitions are not \n",
    "actually implemented in the Lambda function. We are only adding them to demonstrate a large volume of tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2599c9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "calc_api_spec = read_apispec(\"./calc/calc-api.json\")\n",
    "print(f\"API spec for calc has {len(calc_api_spec)} functions\\n\")\n",
    "calc_lambda_arn = calc_lambda_resp[\"lambda_function_arn\"]\n",
    "print(f\"Calc Lambda ARN: {calc_lambda_arn}\")\n",
    "\n",
    "time.sleep(5)\n",
    "calcTargetId = create_gatewaytarget(\n",
    "    gateway_id=gatewayId,\n",
    "    lambda_arn=calc_lambda_arn,\n",
    "    target_name=\"CalcTools\",\n",
    "    target_descr=\"Calculation Tools\",\n",
    "    api_spec=calc_api_spec,\n",
    ")\n",
    "print(f\"CalcTools Target created with id: {calcTargetId} on gateway: {gatewayId}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3200dc9f",
   "metadata": {},
   "source": [
    "To demonstrate the power of gateway search, now we add a few more copies of the Calculator target, \n",
    "so that we end up with 300+ MCP tools exposed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b3aedf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_more_tools(gatewayId):\n",
    "    time.sleep(10)\n",
    "    calcTargetId = create_gatewaytarget(\n",
    "        gateway_id=gatewayId,\n",
    "        lambda_arn=calc_lambda_arn,\n",
    "        target_name=\"Calc2\",\n",
    "        target_descr=\"Calculation 2 Tools\",\n",
    "        api_spec=calc_api_spec,\n",
    "    )\n",
    "    print(f\"Calc2 Target created with id: {calcTargetId} on gateway: {gatewayId}\")\n",
    "    time.sleep(10)\n",
    "    calcTargetId = create_gatewaytarget(\n",
    "        gateway_id=gatewayId,\n",
    "        lambda_arn=calc_lambda_arn,\n",
    "        target_name=\"Calc3\",\n",
    "        target_descr=\"Calculation 3 Tools\",\n",
    "        api_spec=calc_api_spec,\n",
    "    )\n",
    "    print(f\"Calc3 Target created with id: {calcTargetId} on gateway: {gatewayId}\")\n",
    "    time.sleep(10)\n",
    "    calcTargetId = create_gatewaytarget(\n",
    "        gateway_id=gatewayId,\n",
    "        lambda_arn=calc_lambda_arn,\n",
    "        target_name=\"Calc4\",\n",
    "        target_descr=\"Calculation 4 Tools\",\n",
    "        api_spec=calc_api_spec,\n",
    "    )\n",
    "    print(f\"Calc4 Target created with id: {calcTargetId} on gateway: {gatewayId}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62852ae9",
   "metadata": {},
   "outputs": [],
   "source": [
    "add_more_tools(gatewayId=gatewayId)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3717e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = agentcore_client.list_gateway_targets(gatewayIdentifier=gatewayId)\n",
    "targets = resp[\"items\"]\n",
    "for target in resp[\"items\"]:\n",
    "    print(f\"{target['name']} - {target['description']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b701bf49",
   "metadata": {},
   "source": [
    "# Searching for tools from a Gateway"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5599870",
   "metadata": {},
   "source": [
    "### Getting familiar with MCP list tools before we search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6dbea12",
   "metadata": {},
   "source": [
    "Let's define some utility functions to retrieve our MCP endpoint URL for a given Gateway ID, and to \n",
    "retrieve our JWT OAuth access token to securely use our Gateway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d15ac0f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_gateway_endpoint(gateway_id):\n",
    "    response = agentcore_client.get_gateway(gatewayIdentifier=gateway_id)\n",
    "    gateway_url = response[\"gatewayUrl\"]\n",
    "    return gateway_url"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cac832e",
   "metadata": {},
   "source": [
    "Now that our Gateway is created and has targets, let's grab the MCP URL to that\n",
    "Gateway. We can retrieve the endpoint URL from the Gateway control plane based on the Gateway ID."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96c4708f",
   "metadata": {},
   "source": [
    "#### Using MCP Inspector against your Gateway\n",
    "\n",
    "Now that we have an endpoint URL for the MCP server, and we have a JWT bearer token, you may want to explore\n",
    "the MCP server with the MCP Inspector tool. MCP Inspector is an open source tool that can connect to any MCP\n",
    "server, lets you list the tools provided, and even provides an easy to use tool invocation experience. \n",
    "\n",
    "From your terminal window, simply enter `npx @modelcontextprotocol/inspector` to launch the MCP Inspector. Then paste\n",
    "in your Gateway endpoint URL and your JWT token to connect. Once connected, try out List Tools and Invoke Tool.\n",
    "\n",
    "Here's a sample screenshot.\n",
    "\n",
    "![MCP Inspector](images/mcp_inspector.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97f5f769",
   "metadata": {},
   "outputs": [],
   "source": [
    "gatewayEndpoint = get_gateway_endpoint(gateway_id=gatewayId)\n",
    "print(f\"Gateway Endpoint - MCP URL: {gatewayEndpoint}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a99920e9",
   "metadata": {},
   "source": [
    "MCP server security is based on OAuth. To interact with our Gateway, we'll need to\n",
    "retrieve a JWT OAuth access token from our IdP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb2525ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "jwtToken = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")\n",
    "print(f\"Bearer token: {jwtToken}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68bcd0e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !npx @modelcontextprotocol/inspector"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de5f7d58",
   "metadata": {},
   "source": [
    "#### Creating helper functions that use jsonrpc to invoke MCP tools or list them\n",
    "Let's define a helper function called `invoke_gateway_tool` that uses jsonrpc to invoke any of the\n",
    "MCP tools exposed by an MCP Server, including of course, your Gateway. Given an endpoint URL and a JWT token,\n",
    "you can use this utility to invoke any of the MCP tools that AgentCore Gateway made available\n",
    "for you when you added Gateway Targets to your Gateway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "430730db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def invoke_gateway_tool(gateway_endpoint, jwt_token, tool_params):\n",
    "    # print(f\"Invoking tool {tool_params['name']}\")\n",
    "\n",
    "    requestBody = {\n",
    "        \"jsonrpc\": \"2.0\",\n",
    "        \"id\": 2,\n",
    "        \"method\": \"tools/call\",\n",
    "        \"params\": tool_params,\n",
    "    }\n",
    "    response = requests.post(\n",
    "        gateway_endpoint,\n",
    "        json=requestBody,\n",
    "        headers={\n",
    "            \"Authorization\": f\"Bearer {jwt_token}\",\n",
    "            \"Content-Type\": \"application/json\",\n",
    "        },\n",
    "    )\n",
    "\n",
    "    return response.json()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f38bb52f",
   "metadata": {},
   "source": [
    "Here's another utility function for using MCP's `tools/list` method for listing the MCP tools\n",
    "available from your Gateway. Given a Gateway ID and a JWT Token, it retrieves the full set\n",
    "of tools from that Gateway, and returns a list in agent-ready form. The returned list contains \n",
    "Strands Agents MCPAgentTool objects that are suitable for handing your Agent. \n",
    "\n",
    "Note that `tools/list` call is paginated, so the function needs to loop, getting a page of\n",
    "tools at a time, until the `nextCursor` field is no longer populated. The utility function directly\n",
    "calls the endpoint using HTTPS and the jsonrpc protocol. This is a lower level way to list tools\n",
    "compared to the `MCPClient` class provided by Strands Agents. We'll see that experience later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d21a000",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_all_agent_tools_from_mcp_endpoint(gateway_endpoint, jwt_token, client):\n",
    "    more_tools = True\n",
    "    tools_count = 0\n",
    "    tools_list = []\n",
    "\n",
    "    requestBody = {\"jsonrpc\": \"2.0\", \"id\": 2, \"method\": \"tools/list\", \"params\": {}}\n",
    "    next_cursor = \"\"\n",
    "\n",
    "    while more_tools:\n",
    "        if tools_count == 0:\n",
    "            requestBody[\"params\"] = {}\n",
    "        else:\n",
    "            print(f\"\\nGetting next page of tools since a next cursor was returned\\n\")\n",
    "            requestBody[\"params\"] = {\"cursor\": next_cursor}\n",
    "\n",
    "        headers = {\n",
    "            \"Authorization\": f\"Bearer {jwt_token}\",\n",
    "            \"Content-Type\": \"application/json\",\n",
    "        }\n",
    "\n",
    "        print(f\"\\n\\nListing tools for gateway {gateway_endpoint}\")\n",
    "\n",
    "        response = requests.post(gateway_endpoint, json=requestBody, headers=headers)\n",
    "\n",
    "        tools_json = response.json()\n",
    "        tools_count += len(tools_json[\"result\"][\"tools\"])\n",
    "\n",
    "        for tool in tools_json[\"result\"][\"tools\"]:\n",
    "            mcp_tool = MCPTool(\n",
    "                name=tool[\"name\"],\n",
    "                description=tool[\"description\"],\n",
    "                inputSchema=tool[\"inputSchema\"],\n",
    "            )\n",
    "            mcp_agent_tool = MCPAgentTool(mcp_tool, client)\n",
    "            short_descr = tool[\"description\"][0:40] + \"...\"\n",
    "            print(f\"adding tool '{mcp_agent_tool.tool_name}' - {short_descr}\")\n",
    "            tools_list.append(mcp_agent_tool)\n",
    "\n",
    "        if \"nextCursor\" in tools_json[\"result\"]:\n",
    "            next_cursor = tools_json[\"result\"][\"nextCursor\"]\n",
    "            more_tools = True\n",
    "        else:\n",
    "            more_tools = False\n",
    "\n",
    "    print(f\"\\nTotal tools found: {tools_count}\\n\")\n",
    "    return tools_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29ea36db",
   "metadata": {},
   "source": [
    "Lets use this helper function and see the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3eff51d",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")\n",
    "with client:\n",
    "    all_tools = get_all_agent_tools_from_mcp_endpoint(\n",
    "        gateway_endpoint=gatewayEndpoint, jwt_token=jwtToken, client=client\n",
    "    )\n",
    "    print(f\"\\nFound {len(all_tools)} tools using jsonrpc to list MCP tools\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "271d04ef",
   "metadata": {},
   "source": [
    "#### Using Strands Agents list_tools_sync() with pagination\n",
    "If you have written any Python based MCP client, you are likely familiar with the `list_tools_sync()` method \n",
    "which returns the set of tools available from the MCP Server which the client is associated\n",
    "with. But did you know MCP list tools is also paginated? By default, you will only get the first small\n",
    "subset of tools returned. For simple MCP servers, you may not have noticed this, but for many real world \n",
    "MCP servers, your code needs to loop, grabbing pages of tools at a time\n",
    "until there are no more pages remaining. The following utility `get_all_mcp_tools_from_mcp_client` does exactly that. \n",
    "It returns the full list of tools from a given Strands Agent MCP Client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef10b8c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_all_mcp_tools_from_mcp_client(client):\n",
    "    more_tools = True\n",
    "    tools = []\n",
    "    pagination_token = None\n",
    "    while more_tools:\n",
    "        tmp_tools = client.list_tools_sync(pagination_token=pagination_token)\n",
    "        tools.extend(tmp_tools)\n",
    "        if tmp_tools.pagination_token is None:\n",
    "            more_tools = False\n",
    "        else:\n",
    "            more_tools = True\n",
    "            pagination_token = tmp_tools.pagination_token\n",
    "    return tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e519196b",
   "metadata": {},
   "source": [
    "Let's give it a try with our Gateway and find out how many tools the Python client finds. \n",
    "First we create an MCPClient object based on our endpoint URL and our JWT bearer token. Then \n",
    "we retrieve the full set of tools across many pages of tools returned by the MCP server.\n",
    "Given the targets we added earlier, this should return 300+ tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "817d1e48",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")\n",
    "with client:\n",
    "    all_tools = get_all_mcp_tools_from_mcp_client(client)\n",
    "    print(f\"\\nFound {len(all_tools)} tools from list_tools_sync() on mcp client\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba410062",
   "metadata": {},
   "source": [
    "We have now seen 3 different ways to get the full set of tools from your Gateway using it\n",
    "as an MCP Server: \n",
    "\n",
    "1. directly using jsonrpc\n",
    "2. using the `list_tools_sync()` method on the Strands Agent MCPClient\n",
    "3. using the MCP Inspector tool (which uses jsonrpc behind the scenes). \n",
    "\n",
    "For typical developers building agents, you'll be using option 2."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af97ef22",
   "metadata": {},
   "source": [
    "### Using the built-in Gateway semantic search tool\n",
    "Now lets try our first semantic search on the Gateway using its built-in search tool provided as\n",
    "an additional MCP tool that gets added to your MCP tool list."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9de8a698",
   "metadata": {},
   "source": [
    "First let's define a simple utility function to execute the search tool using MCP.\n",
    "Just like for listing tools, we need the gateway endpoint and JWT token. Other than that,\n",
    "all we need to pass in is the search query. The Gateway search tool will do the rest,\n",
    "matching that query against the serverless vector store that it automatically manages on your behalf."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04d868a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tool_search(gateway_endpoint, jwt_token, query):\n",
    "    toolParams = {\n",
    "        \"name\": \"x_amz_bedrock_agentcore_search\",\n",
    "        \"arguments\": {\"query\": query},\n",
    "    }\n",
    "    toolResp = invoke_gateway_tool(\n",
    "        gateway_endpoint=gateway_endpoint, jwt_token=jwt_token, tool_params=toolParams\n",
    "    )\n",
    "    tools = toolResp[\"result\"][\"structuredContent\"][\"tools\"]\n",
    "    return tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "939755a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "tools_found = tool_search(\n",
    "    gateway_endpoint=gatewayEndpoint,\n",
    "    jwt_token=jwtToken,\n",
    "    query=\"find me 3 credit research tools\",\n",
    ")\n",
    "end_time = time.time()\n",
    "print(\n",
    "    f\"tool search via direct Gateway invocation took {(end_time - start_time):.2f} seconds\"\n",
    ")\n",
    "print(f\"Top tool: {tools_found[0]['name']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "602cbd2d",
   "metadata": {},
   "source": [
    "Notice how fast the search returns, in under a second in most cases. The results are returned\n",
    "in descending order of search relevance based on matching the query to the tool metadata.\n",
    "The most relevant tools are first on the list. The intial implementation of search gives back\n",
    "up to 10 results. You could then use all of these tools in your agent, or simply pick a subset of\n",
    "the most relevant matches."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cc03391",
   "metadata": {},
   "source": [
    "# Using Strands Agents with an MCP server that has many tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "433d0f3f",
   "metadata": {},
   "source": [
    "First, we select a model to use with our Strands Agent. \n",
    "For this notebook, we are using Amazon Bedrock models, but Strands and AgentCore\n",
    "can work with any LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b18c10bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrockmodel = BedrockModel(\n",
    "    model_id=\"us.anthropic.claude-3-7-sonnet-20250219-v1:0\",\n",
    "    temperature=0.7,\n",
    "    streaming=True,\n",
    "    boto_session=session,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee8a436a",
   "metadata": {},
   "source": [
    "#### Simple Strands Agent using AgentCore Gateway for agent tools\n",
    "Now lets show how easy it is to use a Strands Agent to leverage an MCP Server\n",
    "provided by AgentCore Gateway. In our\n",
    "simple example, we ask the agent to add some numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2086a0d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "jwtToken = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")\n",
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")\n",
    "with client:\n",
    "    all_tools = get_all_mcp_tools_from_mcp_client(client)\n",
    "    print(f\"\\nFound {len(all_tools)} tools from list_tools_sync() on mcp client\\n\")\n",
    "\n",
    "    simple_agent = Agent(\n",
    "        model=bedrockmodel, tools=all_tools, callback_handler=null_callback_handler\n",
    "    )\n",
    "    result = simple_agent(\"add 100 plus 50 pass \")\n",
    "    print(f\"{result.message['content'][0]['text']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd8d6b33",
   "metadata": {},
   "source": [
    "The Strands Agents framework also lets you bypass the agent event loop, invoking an MCP tool directly.\n",
    "Since Gateway tools are exposed as native MCP tools, this can be done against Gateway tools as well. Here\n",
    "we call a Gateway MCP tool using the `agent.tool.<tool_name>(args)` syntax:\n",
    "\n",
    "```python\n",
    "direct_result = simple_agent.tool.Calc2___add_numbers(firstNumber=10, secondNumber=20)\n",
    "resp_json = json.loads(direct_result['content'][0]['text'])\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddda7a89",
   "metadata": {},
   "outputs": [],
   "source": [
    "jwtToken = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")\n",
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")\n",
    "with client:\n",
    "    all_tools = get_all_mcp_tools_from_mcp_client(client)\n",
    "    print(f\"\\nFound {len(all_tools)} tools from list_tools_sync() on mcp client\\n\")\n",
    "\n",
    "    simple_agent = Agent(\n",
    "        model=bedrockmodel, tools=all_tools, callback_handler=null_callback_handler\n",
    "    )\n",
    "    direct_result = simple_agent.tool.Calc2___add_numbers(\n",
    "        firstNumber=10, secondNumber=20\n",
    "    )\n",
    "    print(f\"direct result = {direct_result}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92d3f963",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_search_tool(client):\n",
    "    mcp_tool = MCPTool(\n",
    "        name=\"x_amz_bedrock_agentcore_search\",\n",
    "        description=\"A special tool that returns a trimmed down list of tools given a context. Use this tool only when there are many tools available and you want to get a subset that matches the provided context.\",\n",
    "        inputSchema={\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "                \"query\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"search query to use for finding tools\",\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"query\"],\n",
    "        },\n",
    "    )\n",
    "    return MCPAgentTool(mcp_tool, client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c04e47cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_using_strands(client, query):\n",
    "    simple_agent = Agent(\n",
    "        model=bedrockmodel,\n",
    "        tools=[get_search_tool(client)],\n",
    "        callback_handler=null_callback_handler,\n",
    "    )\n",
    "\n",
    "    direct_result = simple_agent.tool.x_amz_bedrock_agentcore_search(query=query)\n",
    "\n",
    "    resp_json = json.loads(direct_result[\"content\"][0][\"text\"])\n",
    "    search_results = resp_json[\"tools\"]\n",
    "    # print(json.dumps(search_results, indent=4))\n",
    "    return search_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f07719a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_strands_tools(client, query, top_n):\n",
    "    strands_mcp_tools = []\n",
    "    results = search_using_strands(client, query)\n",
    "    for tool in results[:top_n]:\n",
    "        mcp_tool = MCPTool(\n",
    "            name=tool[\"name\"],\n",
    "            description=tool[\"description\"],\n",
    "            inputSchema=tool[\"inputSchema\"],\n",
    "        )\n",
    "        strands_mcp_tools.append(MCPAgentTool(mcp_tool, client))\n",
    "    return strands_mcp_tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc40ff86",
   "metadata": {},
   "outputs": [],
   "source": [
    "jwtToken = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")\n",
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")\n",
    "with client:\n",
    "    simple_agent = Agent(\n",
    "        model=bedrockmodel,\n",
    "        tools=[get_search_tool(client)],\n",
    "        callback_handler=null_callback_handler,\n",
    "    )\n",
    "\n",
    "    direct_result = simple_agent.tool.x_amz_bedrock_agentcore_search(\n",
    "        query=\"find equity trading tools\"\n",
    "    )\n",
    "\n",
    "    resp_json = json.loads(direct_result[\"content\"][0][\"text\"])\n",
    "    search_results = resp_json[\"tools\"]\n",
    "    print(json.dumps(search_results, indent=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09bd76c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "jwtToken = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")\n",
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")\n",
    "with client:\n",
    "    results = search_using_strands(client, \"find trading tools\")\n",
    "    print(json.dumps(search_results[0], indent=4))\n",
    "\n",
    "    results = search_using_strands(client, \"find credit research tools\")\n",
    "    print(json.dumps(search_results[0], indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c0dcc80",
   "metadata": {},
   "source": [
    "# Adding tool search results to a Strands Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "374bc324",
   "metadata": {},
   "source": [
    "Now let's look at how the tools returned from a search can be added to a\n",
    "Strands Agent. To make the coding simpler, let's provide a utility function that\n",
    "maps tool search results to Strands MCPAgentTool objects. Simply pass in the \n",
    "search results, and indicate how many of those results you want to pass to your\n",
    "agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df1e3efe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tools_to_strands_mcp_tools(tools, top_n):\n",
    "    strands_mcp_tools = []\n",
    "    for tool in tools[:top_n]:\n",
    "        mcp_tool = MCPTool(\n",
    "            name=tool[\"name\"],\n",
    "            description=tool[\"description\"],\n",
    "            inputSchema=tool[\"inputSchema\"],\n",
    "        )\n",
    "        strands_mcp_tools.append(MCPAgentTool(mcp_tool, client))\n",
    "    return strands_mcp_tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d35be1",
   "metadata": {},
   "outputs": [],
   "source": [
    "jwtToken = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")\n",
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")\n",
    "with client:\n",
    "    agent = Agent(\n",
    "        model=bedrockmodel,\n",
    "        tools=find_strands_tools(\n",
    "            client,\n",
    "            \"tools for doing addition, subtraction, multiplication, division\",\n",
    "            10,\n",
    "        ),\n",
    "    )\n",
    "    result = agent(\"(10*2)/(5-3)\")\n",
    "    print(f\"{result.message['content'][0]['text']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6de85475",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "jwtToken = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")\n",
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")\n",
    "with client:\n",
    "    print(\"Searching for an ADDING tool from endpoint with full set of tools...\")\n",
    "    tools_found = tool_search(\n",
    "        gateway_endpoint=gatewayEndpoint,\n",
    "        jwt_token=jwtToken,\n",
    "        query=\"tools for multiplying two numbers\",\n",
    "    )\n",
    "    print(f\"Top tool found: {tools_found[0]['name']}\\n\")\n",
    "\n",
    "    agent = Agent(model=bedrockmodel, tools=tools_to_strands_mcp_tools(tools_found, 1))\n",
    "    result = agent(\"10 * 70\")\n",
    "    print(f\"{result.message['content'][0]['text']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a41cfa0",
   "metadata": {},
   "source": [
    "Notice the latency improvement. This example using a subset of tools from Gateway search is significantly faster than\n",
    "agent invocation when depending on hundreds of tools."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78c1e64e",
   "metadata": {},
   "source": [
    "# Showing 3x latency improvement by using tool search"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bc164be",
   "metadata": {},
   "source": [
    "Now that we know how to use Gateway MCP tools from a Strands agent, and we know how to search for tools and\n",
    "add them to an agent, lets show the power of search. We'll highlight the significant latency reduction\n",
    "and input token usage that can be delivered.\n",
    "\n",
    "To demonstrate the latency and token reductions, we compare 2 approaches side by side:\n",
    "\n",
    "1. **Without search**. We add the full set of MCP tools that the MCP server exposes (300+ in our case) to our agent and let the agent do its tool selection and invocation accordingly.\n",
    "2. **Using search**. In the second approach, we do a search based on the topic at hand, and only send in the most relevant tools to the agent. To prove the point, we use two different topics: math (adding numbers), and food (booking a restaurant reservation), each requiring a different set of tools.\n",
    "\n",
    "To normalize the latency distribution and get a meaningful comparison, we perform multiple iterations of\n",
    "each approach. Also, to avoid overstating the gains, when doing the search approach, we include not only the\n",
    "latency of the agent invocation, but also the latency of performing the tool search. For each \n",
    "iteration, we hand the agent two tasks: \n",
    "\n",
    "1. Math task -  add 2 numbers \n",
    "2. Food task - book a restaurant reservation\n",
    "\n",
    "The results below demonstrate the benefits, highlighting 3x latency reduction, and even greater reduction in\n",
    "input token usage. Note that while token usage savings translate to cost savings, that may not be as impactful due to\n",
    "the relatively lower cost of input tokens (for many model providers, input tokens are much lest costly). Even so, for \n",
    "large scale agent deployment, even input token usage costs can add up, so dynamic search can help reduce \n",
    "agent runtime costs as well."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e24ad450",
   "metadata": {},
   "source": [
    "#### Measure latency and token usage for agent using the entire set of MCP tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e67f2ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "iterations = 2\n",
    "full_tokens = light_tokens = 0\n",
    "full_elapsed_time = light_elapsed_time = 0\n",
    "\n",
    "jwtToken = utils.get_bearer_token(\n",
    "    client_id=cognito_response[\"client_id\"],\n",
    "    username=\"testuser\",\n",
    "    password=\"MyPassword123!\",\n",
    ")\n",
    "client = MCPClient(\n",
    "    lambda: streamablehttp_client(\n",
    "        f\"{gatewayEndpoint}\", headers={\"Authorization\": f\"Bearer {jwtToken}\"}\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bdf1296",
   "metadata": {},
   "outputs": [],
   "source": [
    "with client:\n",
    "    all_tools = get_all_mcp_tools_from_mcp_client(client)\n",
    "    print(f\"\\nFound {len(all_tools)} tools from list_tools_sync() on mcp client\\n\")\n",
    "    heavy_agent = Agent(\n",
    "        model=bedrockmodel, tools=all_tools, callback_handler=null_callback_handler\n",
    "    )\n",
    "\n",
    "    math_input = \"add 100 plus <iteration>\"\n",
    "    food_input = (\n",
    "        \"book me a table for 2 at Burger King under name Jo Smith at 7pm August <day>\"\n",
    "    )\n",
    "\n",
    "    print(\"using agent with ALL tools...\")\n",
    "    start_time = time.time()\n",
    "\n",
    "    for i in range(iterations):\n",
    "        result = heavy_agent(math_input.replace(\"<iteration>\", str(i + 1)))\n",
    "        print(f\"{i+1}) {result.message['content'][0]['text']}\")\n",
    "\n",
    "        result = heavy_agent(food_input.replace(\"<day>\", str(i + 1)))\n",
    "        print(f\"{i+1}) {result.message['content'][0]['text']}\")\n",
    "\n",
    "    end_time = time.time()\n",
    "    full_tokens = result.metrics.accumulated_usage[\"totalTokens\"]\n",
    "    full_elapsed_time = end_time - start_time\n",
    "    print(f\"\\nTotal time: {full_elapsed_time:.1f} s, tokens: {full_tokens:,d}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be6b20a3",
   "metadata": {},
   "source": [
    "#### Measure latency and token usage for agents using Gateway Search\n",
    "Now we'll use a dynamic approach, calling search to find relevant tools, and then calling the\n",
    "agent with only those relevant tools. Note that since we are resetting the agent on each \n",
    "conversation turn, we're also intializing the message list from conversation history of the prior turn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80296fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "with client:\n",
    "    print(\"using agent with ONLY tools from focused search...\")\n",
    "    start_time = time.time()\n",
    "    messages = []\n",
    "\n",
    "    light_agent = Agent()\n",
    "\n",
    "    for i in range(iterations):\n",
    "        print(\"Searching for an ADDING tool from endpoint with full set of tools...\")\n",
    "        tools_found = tool_search(\n",
    "            gateway_endpoint=gatewayEndpoint,\n",
    "            jwt_token=jwtToken,\n",
    "            query=\"tools for simply adding two numbers\",\n",
    "        )\n",
    "        print(f\"Top tool found: {tools_found[0]['name']}\\n\")\n",
    "        light_agent = Agent(\n",
    "            model=bedrockmodel,\n",
    "            tools=tools_to_strands_mcp_tools(tools_found, 1),\n",
    "            messages=messages,\n",
    "            callback_handler=null_callback_handler,\n",
    "        )\n",
    "        light_result = light_agent(math_input.replace(\"<iteration>\", str(i + 1)))\n",
    "        print(f\"{i+1}) {light_result.message['content'][0]['text']}\")\n",
    "        messages = light_agent.messages\n",
    "\n",
    "        print(\n",
    "            \"Searching for a RESTAURANT BOOKING tool from endpoint with full set of tools...\"\n",
    "        )\n",
    "        tools_found = tool_search(\n",
    "            gateway_endpoint=gatewayEndpoint,\n",
    "            jwt_token=jwtToken,\n",
    "            query=\"tools for booking a restaurant reservation\",\n",
    "        )\n",
    "        print(f\"Top tool found: {tools_found[0]['name']}\\n\")\n",
    "        light_agent = Agent(\n",
    "            model=bedrockmodel,\n",
    "            tools=tools_to_strands_mcp_tools(tools_found, 1),\n",
    "            messages=messages,\n",
    "            callback_handler=null_callback_handler,\n",
    "        )\n",
    "        light_result = light_agent(food_input.replace(\"<day>\", str(i + 1)))\n",
    "        print(f\"{i+1}) {light_result.message['content'][0]['text']}\")\n",
    "        messages = light_agent.messages\n",
    "        light_tokens = light_result.metrics.accumulated_usage[\"totalTokens\"]\n",
    "    end_time = time.time()\n",
    "\n",
    "    light_elapsed_time = end_time - start_time\n",
    "    print(f\"\\nTotal time: {light_elapsed_time:.1f} s, tokens: {light_tokens:,d}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2647422",
   "metadata": {},
   "source": [
    "#### Compare results, higlighting benefits of search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef22580",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    f\"\\n\\nLatency without search: {full_elapsed_time:.1f}s, using search: {light_elapsed_time:.1f}s\"\n",
    ")\n",
    "print(f\"Tokens without search: {full_tokens:,d}, using search: {light_tokens:,d}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7a0d35e",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "In this tutorial, you have learned about Amazon Bedrock AgentCore Gateway and its built-in \n",
    "fully managed semantic search capability. You have seen the following:\n",
    "\n",
    "- how to create a gateway with semantic search enabled\n",
    "- how to add multiple gateway targets to surface 300+ MCP tools from a single endpoint\n",
    "- how to list the tools on your gateway using 3 different approaches\n",
    "- how to use the built-in semantic search tool to find relevant tools\n",
    "- how to integrate search with your Strands Agent\n",
    "- how to compare performance of an agent using a server with hundreds of tools versus one that uses semantic search to narrow tools to a specific topic\n",
    "\n",
    "AgentCore Gateway search is helpful for more advanced use cases as well. By offering the search as a native\n",
    "MCP tool and not just a control plane API, you can imagine giving your agents more autonomy to discover new\n",
    "MCP servers, and find new capabilities at runtime leading to breakthroughs in solving more challenging problems.\n",
    "In addition, search is an important foundation for MCP registries and supporting agent developers as they \n",
    "design and build new agents."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c569fc6",
   "metadata": {},
   "source": [
    "# Cleaning up resources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c51ae00",
   "metadata": {},
   "source": [
    "First let's define some helper functions for cleaning up AgentCore Gateway resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81ce65be",
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_gatewaytarget(gateway_id):\n",
    "    response = agentcore_client.list_gateway_targets(gatewayIdentifier=gateway_id)\n",
    "\n",
    "    print(f\"Found {len(response['items'])} targets for the gateway\")\n",
    "\n",
    "    for target in response[\"items\"]:\n",
    "        print(\n",
    "            f\"Deleting target with Name: {target['name']} and Id: {target['targetId']}\"\n",
    "        )\n",
    "\n",
    "        response = agentcore_client.delete_gateway_target(\n",
    "            gatewayIdentifier=gateway_id, targetId=target[\"targetId\"]\n",
    "        )\n",
    "        time.sleep(20)\n",
    "\n",
    "\n",
    "def delete_gateway(gateway_id):\n",
    "    response = agentcore_client.delete_gateway(gatewayIdentifier=gateway_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecaade0e",
   "metadata": {},
   "source": [
    "### Deleting Gateway Targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "930b0d93",
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_gatewaytarget(gateway_id=gatewayId)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77c008ef",
   "metadata": {},
   "source": [
    "### Deleting the Gateway itself"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78916cc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_gateway(gateway_id=gatewayId)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa6fc55e",
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_arns = [\n",
    "    calc_lambda_resp[\"lambda_function_arn\"],\n",
    "    restaurant_lambda_resp[\"lambda_function_arn\"],\n",
    "]\n",
    "\n",
    "for arn in lambda_arns:\n",
    "    if utils.delete_gateway_lambda(arn):\n",
    "        print(f\"Deleted Lambda: {arn}\")\n",
    "    else:\n",
    "        print(f\"Lambda {arn} not found or deletion failed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f507063",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gateway role cleanup\n",
    "if utils.delete_gateway_iam_role():\n",
    "    print(\"Gateway IAM role deleted\")\n",
    "else:\n",
    "    print(\"Gateway IAM role not found or deletion failed\")\n",
    "\n",
    "# Cognito cleanup\n",
    "if utils.delete_cognito_user_pool():\n",
    "    print(\"Cognito pool deleted\")\n",
    "else:\n",
    "    print(\"✗ Failed to delete Cognito pool\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}