{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import stackprinter  # type: ignore\n",
    "import jupyter_black  # type: ignore\n",
    "from dotenv import load_dotenv  # type: ignore\n",
    "\n",
    "from baml_agents import init_logging, with_model\n",
    "from baml_client.sync_client import b\n",
    "from baml_client.type_builder import TypeBuilder as TB\n",
    "from baml_client import types as T\n",
    "\n",
    "init_logging()\n",
    "stackprinter.set_excepthook()\n",
    "load_dotenv()\n",
    "jupyter_black.load()\n",
    "\n",
    "b = with_model(b, \"gpt-4.1-nano\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Model Context Protocol (MCP) Tools with baml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is Model Context Protocol (MCP)\n",
    "\n",
    "Model Context Protocol (MCP) allows users to easily plug-and-play 3rd party resources into AI systems with an interface that's easy to understand and use by LLMs.\n",
    "\n",
    "- **Tip** You can find a list of more popular Model Context Protocol (MCP) resources at [punkpeye/awesome-mcp-servers](https://github.com/punkpeye/awesome-mcp-servers)\n",
    "\n",
    "In this example, we'll focus on integrating MCP Tools into the baml structured generation framework, since MCP offers additional features beyond basic tool invocation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "\n",
    "We'll use the `mcpt` command-line interface to interact with MCP servers. It functions much like `curl`, allowing us to query the servers easily.\n",
    "\n",
    "- Install it from here: [github.com/f/mcptools](https://github.com/f/mcptools)\n",
    "\n",
    "It's not an efficient way because we start and stop the server on every request, but it's a good way to test the integration. In a production environment, you would run the server continuously to always be available for requests.\n",
    "\n",
    "Note: While this example uses local servers (via stdio), you can also specify any URL as the server."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Showing around the MCP Server\n",
    "\n",
    "Let's see all available tools on the `uvx mcp-server-calculator` [MCP server](https://github.com/githejie/mcp-server-calculator). Server will be downloaded as a python package by `uvx` and run locally until we get the data we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[36mcalculate\u001b[0m(\u001b[32mexpression:str\u001b[0m)\n",
      "     \u001b[37mCalculates/evaluates the given expression.\u001b[0m\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Show all the tools in the server\n",
    "!mcpt tools uvx mcp-server-calculator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the `calculate` tool:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[37m246.8\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# Use the calculator tool\n",
    "!mcpt call calculate -p '{\"expression\": \"123.4 * 2\"}' uvx mcp-server-calculator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calling tools from the MCP server with BAML"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see the actual prompt that we will send to the LLM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[system]\n",
      "Please select the next best action to achieve the goal.\n",
      "\n",
      "‹goal›\n",
      "What's the sum of the first 5 prime numbers?\n",
      "‹/goal›\n",
      "\n",
      "Answer in JSON format with exactly one next best action:\n",
      "{\n",
      "  chosen_action: {\n",
      "    // Calculates/evaluates the given expression.\n",
      "    action_id: \"calculate\",\n",
      "    expression: string,\n",
      "  },\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "from baml_agents import ActionRunner, display_prompt\n",
    "\n",
    "calculator_mcp = \"uvx mcp-server-calculator\"\n",
    "r = ActionRunner(TB, b=b, cache=True)\n",
    "r.add_from_mcp_server(calculator_mcp)\n",
    "\n",
    "question = \"What's the sum of the first 5 prime numbers?\"\n",
    "request = r.b.request.BamlCustomTools_GetNextAction(question)\n",
    "display_prompt(request)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's send it and run the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "chosen_action={'action_id': 'calculate', 'expression': '2 + 3 + 5 + 7 + 11'}\n"
     ]
    }
   ],
   "source": [
    "calculator_mcp = \"uvx mcp-server-calculator\"\n",
    "r = ActionRunner(TB, b=b, cache=True)\n",
    "r.add_from_mcp_server(calculator_mcp)\n",
    "\n",
    "question = \"What's the sum of the first 5 prime numbers?\"\n",
    "action = r.b.BamlCustomTools_GetNextAction(question)\n",
    "print(action)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Result(content='28', error=False)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r.run(action)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to choose which tools to include into the prompt?\n",
    "\n",
    "You can filter when adding tools to the ActionRunner and also right before sending the prompt to the LLM."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using all the tools in the server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['brave_web_search', 'brave_local_search']\n"
     ]
    }
   ],
   "source": [
    "brave_search_mcp = \"docker run -i --rm -e BRAVE_API_KEY mcp/brave-search\"\n",
    "env = {\"BRAVE_API_KEY\": \"not-needed\"}\n",
    "\n",
    "r = ActionRunner(TB, b=b, cache=True)\n",
    "r.add_from_mcp_server(brave_search_mcp, env=env)\n",
    "print([a.name for a in r.actions])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's say we only want to import the \"brave_web_search\" tool:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['brave_web_search']\n"
     ]
    }
   ],
   "source": [
    "from baml_agents import McpToolDefinition\n",
    "\n",
    "\n",
    "def include(m: McpToolDefinition):\n",
    "    if m.name == \"brave_web_search\":\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "r = ActionRunner(TB, b=b, cache=True)\n",
    "r.add_from_mcp_server(brave_search_mcp, include=include, env=env)\n",
    "print([a.name for a in r.actions])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's say we imported many tools but only want to make one available to the LLM (we'll the use **include()** function):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['brave_web_search', 'brave_local_search']\n"
     ]
    }
   ],
   "source": [
    "r = ActionRunner(TB, b=b, cache=True)\n",
    "r.add_from_mcp_server(brave_search_mcp, env=env)\n",
    "print([a.name for a in r.actions])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[system]\n",
      "Please select the next best action to achieve the goal.\n",
      "\n",
      "‹goal›\n",
      "Find me some sushi\n",
      "‹/goal›\n",
      "\n",
      "Answer in JSON format with exactly one next best action:\n",
      "{\n",
      "  chosen_action: {\n",
      "    // Searches for local businesses and places using Brave's Local Search API. Best for queries related to physical locations, businesses, restaurants, services, etc. Returns detailed information including:\n",
      "    // - Business names and addresses\n",
      "    // - Ratings and review counts\n",
      "    // - Phone numbers and opening hours\n",
      "    // Use this when the query implies 'near me' or mentions specific locations. Automatically falls back to web search if no local results are found.\n",
      "    action_id: \"brave_local_search\",\n",
      "    // Number of results (1-20, default 5)\n",
      "    count: float OR null,\n",
      "    // Local search query (e.g. 'pizza near Central Park')\n",
      "    query: string,\n",
      "  },\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "from baml_client.type_builder import TypeBuilder\n",
    "\n",
    "\n",
    "question = \"Find me some sushi\"\n",
    "\n",
    "\n",
    "def include(m: McpToolDefinition):\n",
    "    if m.name == \"brave_local_search\":\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "request = r.b_(include=include).request.BamlCustomTools_GetNextAction(question)\n",
    "display_prompt(request)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}